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

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

---

前言

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

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

后续可以学习【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可以查看效果：

到这里读者将代码迁移到新机器时，可以不再安装pytorch相关依赖就能使用模型的预测功能，这可以极大的减少所依赖环境的大小。

---

总结

尽可能简单、详细的介绍CPU模式下Pytorch模型转ONNX格式的流程，后续介绍GPU版本的格式转化，学习难度只是有略微提升。