基于mxnet的resnet50模型转ONNX部署的问题记录

Posted 2023-04-07

参考技术A 本文记录基于mxnet训练得到的resnet50模型在转为onnx格式部署时发生的两个错误及解决办法

os: ubuntu 16.04
Mxnet : 1.6.0
onnx: 1.6.0
cuda: 10.2
cudnn: 8.0

resnet50中使用的batchnorm层在转换至onnx时报错不支持属性spatial

mxnet源码mxnet/contrib/onnx/mx2onnx/_op_translation.py 359行 注释掉spatial参数

通过netron.app 查看网络结构, 看到第一个batchnorm层(bn_data)的参数bn_data_gamma异常:

在onnx模型中手动修改bn_data层的gamma参数为1.0

Pytorch模型(.pth)转onnx模型(.onnx)

简介
Open Neural Network Exchange（ONNX，开放神经网络交换）格式，是一个用于表示深度学习模型的标准，可使模型在不同框架之间进行转移。

ONNX是一种针对机器学习所设计的开放式的文件格式，用于存储训练好的模型。它使得不同的人工智能框架（如Pytorch, MXNet）可以采用相同格式存储模型数据并交互。 ONNX的规范及代码主要由微软，亚马逊 ，Facebook 和 IBM等公司共同开发，以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习框架有： Caffe2, PyTorch, MXNet，ML.NET，TensorRT 和 Microsoft CNTK，并且 TensorFlow也非官方的支持ONNX。
代码实现( .pth --> .onnx )

import torch
import torch.onnx
from tinynet import tinynet
from conf import settings
import os

def pth_to_onnx(input, checkpoint, onnx_path, input_names=['input'], output_names=['output'], device='cpu'):
    if not onnx_path.endswith('.onnx'):
        print('Warning! The onnx model name is not correct,\\
              please give a name that ends with \\'.onnx\\'!')
        return 0

    model = tinynet()
    model.load_state_dict(torch.load(checkpoint))
    model.eval()
    # model.to(device)
    
    torch.onnx.export(model, input, onnx_path, verbose=True, input_names=input_names, output_names=output_names)
    print("Exporting .pth model to onnx model has been successful!")

if __name__ == '__main__':
    os.environ['CUDA_VISIBLE_DEVICES']='2'
    checkpoint = './tinynet.pth'
    onnx_path = './tinynet.onnx'
    input = torch.randn(1, 1, 640, 360)
    # device = torch.device("cuda:2" if torch.cuda.is_available() else 'cpu')
    pth_to_onnx(input, checkpoint, onnx_path)