【深度学习框架格式转化】【GPU】Pytorch模型转ONNX模型格式流程详解【入门】
文章目录
前言
神经网络的模型通常在深度学习框架(PyTorc、TensorFlow和Caffe等)下训练得到，这些特定环境的深度学习框架依赖较多，规模较大，不适合在生产环境中安装，onnx支持大多数框架下模型的转换，便于整合模型，并且深度学习模型需要大量的算力才能满足实时运行需求，需要优化模型的运行效率，onnx并则能带来稳定的提速。
onnx还能再转化成TensorRT(GPU)格式和OpenVINO(CPU)格式进行推理，进一步提升速度
后续可以学习【GPU】Pytorch模型转ONNX格式流程详解
PyTorch模型环境搭建(CPU)
博主以伪装对象分割(COS)之PFNet算法为例进行详解：【PFNet-pytorch代码】。
用PyTorch运行一个伪装对象分割模型PFNet，并把模型部署到ONNX Runtime这个推理引擎上。
博主在win10环境下装anaconda环境，搭建PFNet模型运行的PyTorch环境(官网下载地址)
# 创建虚拟环境
conda create -n pytorch2onnx_cpu python=3.10 -y
# 激活环境
activate pytorch2onnx_cpu
# 下载githup源代码到合适文件夹,并cd到代码文件夹内(科学上网)
git clone https://github.com/Mhaiyang/CVPR2021_PFNet.git
# 安装pytorch(cpu)
pip3 install torch torchvision torchaudio
博主在这里不会详细讲解代码内容，只关注代码的使用，即代码的测试过程。源码作者提供了预训练权重和测试数据，博主整理到了【百度云，提取码：a660】上供大家下载。
下载resnet50-19c8e357.pth放置到CVPR2021_PFNet\backbone\resnet下：
下载PFNet.pth放置到CVPR2021_PFNet下：
下载测试数据集CAMO_TestingDataset.zip、CHAMELEON_TestingDataset.zip和COD10K_TestingDataset.zip解压重命名放置到CVPR2021_PFNet\data\test中：
使用预训练权重进行测试，修改infer.py文件内容
# 1.修改infer.py,只保留在test中有的数据集
to_test = OrderedDict([
('CHAMELEON', chameleon_path),
('CAMO', camo_path),
('COD10K', cod10k_path),
# ('NC4K', nc4k_path)
])
# 2.修改infer.py,删除/注释所有使用gpu相关代码
# device_ids = [0]
# torch.cuda.set_device(device_ids[0])
# net = PFNet(backbone_path).cuda(device_ids[0])
net = PFNet(backbone_path)
# img_var = Variable(img_transform(img).unsqueeze(0)).cuda(device_ids[0])
img_var = Variable(img_transform(img).unsqueeze(0))
# 3.修改config.py中的内容
# datasets_root = '../data/NEW'修改成datasets_root = './data
在CVPR2021_PFNet\results可以查看效果：
到这里PyTorch模型环境搭建(CPU)完毕。
安装onnx和onnxruntime(CPU)
需要在anaconda虚拟环境安装onnx和onnxruntime
# 激活环境
activate pytorch2onnx_cpu
# 安装onnx
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple onnx
# 安装CPU版
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple onnxruntime
获取ONNX Runtime的版本信息
import onnxruntime as ort
print("ONNX Runtime version:", ort.__version__)
pytorch2onnx
在CVPR2021_PFNet目录下新建pytorch2onnx.py文件并执行文件
import onnx
from onnx import numpy_helper
import torch
from PFNet import PFNet
backbone_path = './backbone/resnet/resnet50-19c8e357.pth'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
example = torch.randn(1,3, 416, 416).to(device) # 1 3 416 416
print(example.dtype)
model = PFNet(backbone_path) # PFNet网络模型
model.load_state_dict(torch.load(r'PFNet.pth')) # 加载训练好的模型
model = model.to(device) # 模型放到cpu上
model.eval()
torch.onnx.export(model, example, r"PFNet.onnx") # 导出模型
model_onnx = onnx.load(r"PFNet.onnx") # onnx加载保存的onnx模型
onnx.checker.check_model(model_onnx) # 检查模型是否有问题
print(onnx.helper.printable_graph(model_onnx.graph)) # 打印onnx网络
pytorch模型转化成onnx模型成功。
现在抛开任何pytorch相关的依赖，使用onnx模型完成测试，新建run_onnx.py，代码是参考源代码的推理部分infer.py改写来的。
import onnxruntime as ort
import numpy as np
from collections import OrderedDict
from config import *
from PIL import Image
from numpy import mean
import time
import datetime
def composed_transforms(image):
mean = np.array([0.485, 0.456, 0.406]) # 均值
std = np.array([0.229, 0.224, 0.225]) # 标准差
# transforms.Resize是双线性插值
resized_image = image.resize((args['scale'], args['scale']), resample=Image.BILINEAR)
# onnx模型的输入必须是np,并且数据类型与onnx模型要求的数据类型保持一致
resized_image = np.array(resized_image)
normalized_image = (resized_image/255.0 - mean) / std
return np.round(normalized_image.astype(np.float32), 4)
def check_mkdir(dir_name):
if not os.path.exists(dir_name):
os.makedirs(dir_name)
to_test = OrderedDict([
# ('CHAMELEON', chameleon_path),
# ('CAMO', camo_path),
('COD10K', cod10k_path),
])
args = {
'scale': 416,
'save_results': True
}
def main():
# 保存检测结果的地址
results_path = './results2'
exp_name = 'PFNet'
providers = ["CPUxecutionProvider"]
ort_session = ort.InferenceSession("PFNet.onnx", providers=providers) # 创建一个推理session
input_name = ort_session.get_inputs()[0].name
# 输出有四个
output_names = [output.name for output in ort_session.get_outputs()]
start = time.time()
for name, root in to_test.items():
time_list = []
image_path = os.path.join(root, 'image')
if args['save_results']:
check_mkdir(os.path.join(results_path, exp_name, name))
img_list = [os.path.splitext(f)[0] for f in os.listdir(image_path) if f.endswith('jpg')]
for idx, img_name in enumerate(img_list):
img = Image.open(os.path.join(image_path, img_name + '.jpg')).convert('RGB')
w, h = img.size
# 对原始图像resize和归一化
img_var = composed_transforms(img)
# np的shape从[w,h,c]=>[c,w,h]
img_var = np.transpose(img_var, (2, 0, 1))
# 增加数据的维度[c,w,h]=>[bathsize,c,w,h]
img_var = np.expand_dims(img_var, axis=0)
start_each = time.time()
prediction = ort_session.run(output_names, {input_name: img_var})
time_each = time.time() - start_each
time_list.append(time_each)
# 除去多余的bathsize维度,NumPy变会PIL同样需要变换数据类型
# *255替换pytorch的to_pil
prediction = (np.squeeze(prediction[3])*255).astype(np.uint8)
if args['save_results']:
(Image.fromarray(prediction).resize((w, h)).convert('L').save(os.path.join(results_path, exp_name, name, img_name + '.png')))
print(('{}'.format(exp_name)))
print("{}'s average Time Is : {:.3f} s".format(name, mean(time_list)))
print("{}'s average Time Is : {:.1f} fps".format(name, 1 / mean(time_list)))
end = time.time()
print("Total Testing Time: {}".format(str(datetime.timedelta(seconds=int(end - start)))))
if __name__ == '__main__':
main()
在CVPR2021_PFNet\results2可以查看效果：
总结
尽可能简单、详细的介绍CPU模式下Pytorch模型转ONNX格式的流程，后续介绍GPU版本的格式转化，学习难度只是有略微提升。