【深度学习框架格式转化】【CPU】Pytorch模型转ONNX模型格式流程详解

【深度学习框架格式转化】【CPU】Pytorch模型转ONNX模型格式流程详解

提示:博主取舍了很多大佬的博文并亲测有效,分享笔记邀大家共同学习讨论

---

前言

神经网络的模型通常在深度学习框架(PyTorch、TensorFlow和Caffe等)下训练得到，这些特定环境的深度学习框架依赖较多，规模较大，不适合在生产环境中安装，onnx支持大多数框架下模型的转换，便于整合模型，并且深度学习模型需要大量的算力才能满足实时运行需求，需要优化模型的运行效率，而onnx能够带来稳定的提速。

CPU模式下的格式转化，无论Pytorch还是ONNX搭建流程都十分简便，适合入门学习，也对极其适合对硬件要求很低的轻量级模型的运行。

读者可以通过学习【OnnxRuntime部署】系列学习文章目录的快速入门 的内容，快速入门上手。

---

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