Pytorch图像分割模型转ONNX

最近学习了一下图像分割模型部署——PyTorch转ONNX
参考了子豪兄的视频

一、主要学习内容如下：

即ONNX是个中间翻译器，帮助不同框架和不同设备更方便地迁移和部署。

二、具体操作

1. 安装配置环境

安装Pytorch
pip3 install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113
安装 ONNX
pip install onnx -i https://pypi.tuna.tsinghua.edu.cn/simple
安装推理引擎 ONNX Runtime
pip install onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple
验证安装配置成功

import torch
print('PyTorch 版本', torch.__version__)
#PyTorch 版本 1.10.0+cu113
import onnx
print('ONNX 版本', onnx.__version__)
#ONNX 版本 1.13.1
import onnxruntime as ort
print('ONNX Runtime 版本', ort.__version__)
#ONNX Runtime 版本 1.14.1

由于我已经安装过PyTorch和其他第三方包，因此只需要安装ONNX和ONNX Runtime即可。

2. 转ONNX

导入工具包

import torch
from torchvision import models

# 有 GPU 就用 GPU，没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device', device)

导入训练好的模型

model = torch.load('checkpoint/fruit30_pytorch_20220814.pth')
model = model.eval().to(device)

构造一个输入图像Tensor

x = torch.randn(1, 3, 256, 256).to(device)

输入Pytorch模型推理预测，获得1000个类别的预测结果

output = model(x)
print(output.shape)

Pytorch模型转ONNX模型

x = torch.randn(1, 3, 256, 256).to(device)

with torch.no_grad():
    torch.onnx.export(
        model,                   # 要转换的模型
        x,                       # 模型的任意一组输入
        'resnet18_fruit30.onnx', # 导出的 ONNX 文件名
        opset_version=11,        # ONNX 算子集版本
        input_names=['input'],   # 输入 Tensor 的名称（自己起名字）
        output_names=['output']  # 输出 Tensor 的名称（自己起名字）
    ) 

验证onnx模型导出成功

import onnx

# 读取 ONNX 模型
onnx_model = onnx.load('resnet18_fruit30.onnx')

# 检查模型格式是否正确
onnx.checker.check_model(onnx_model)

print('无报错，onnx模型载入成功')

使用Netron可视化模型结构
Netron：https://netron.app
视频教程：https://www.bilibili.com/video/BV1TV4y1P7AP

适用到自己的网络模型：

# ------------------------------------------------------------------------------
# Modified based on https://github.com/HRNet/HRNet-Semantic-Segmentation
# ------------------------------------------------------------------------------

import argparse
import os
import pprint

import logging
import timeit

import numpy as np

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn

import _init_paths
import models
import datasets
from configs import config
from configs import update_config
from utils.function import testval, test, test_ood
from utils.utils import create_logger

# 有 GPU 就用 GPU，没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device', device)

def parse_args():
    parser = argparse.ArgumentParser(description='Train segmentation network')

    parser.add_argument('--cfg',
                        help='experiment configure file name',
                        default="configs/bandon/pidnet_small_bandon.yaml",
                        type=str)
    parser.add_argument('opts',
                        help="Modify config options using the command-line",
                        default=None,
                        nargs=argparse.REMAINDER)

    args = parser.parse_args()
    update_config(config, args)

    return args

def main():
    args = parse_args()

    logger, final_output_dir, _ = create_logger(
        config, args.cfg, 'test')

    logger.info(pprint.pformat(args))
    logger.info(pprint.pformat(config))

    # cudnn related setting
    cudnn.benchmark = config.CUDNN.BENCHMARK
    cudnn.deterministic = config.CUDNN.DETERMINISTIC
    cudnn.enabled = config.CUDNN.ENABLED

    # build model
    model = model = models.pidnet.get_seg_model(config, imgnet_pretrained=True)

    if config.TEST.MODEL_FILE:
        model_state_file = config.TEST.MODEL_FILE
    else:
        model_state_file = os.path.join(final_output_dir, 'best.pt')      

    logger.info('=> loading model from {}'.format(model_state_file))

    pretrained_dict = torch.load(model_state_file)
    if 'state_dict' in pretrained_dict:
        pretrained_dict = pretrained_dict['state_dict']
    model_dict = model.state_dict()
    pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items()
                        if k[6:] in model_dict.keys()}
    for k, _ in pretrained_dict.items():
        logger.info(
            '=> loading {} from pretrained model'.format(k))
    model_dict.update(pretrained_dict)
    model.load_state_dict(model_dict)
    #model = model.cuda()
    model = model.eval().to(device)

    # 包装模型
    #export_model = CustomExportModel(model).to(device)

    x = torch.randn(1, 3, 512, 512).to(device)

    with torch.no_grad():
        torch.onnx.export(
            model,            # 要自定义的模型
            x,                       # 模型的任意一组输入
            'NeSF-Net.onnx',         # 导出的 ONNX 文件名
            opset_version=11,        # ONNX 算子集版本
            input_names=['input'],   # 输入 Tensor 的名称（自己起名字）
            output_names=['output']  # 输出 Tensor 的名称（自己起名字）
        )
    print('finish!')

if __name__ == '__main__':
    main()
    import onnx
    # 读取 ONNX 模型
    onnx_model = onnx.load('NeSF-Net.onnx')

    # 检查模型格式是否正确
    onnx.checker.check_model(onnx_model)

    print('无报错，onnx模型载入成功')