Pytorch模型转ONNX模型

Pytouch模型转ONNX

将PyTorch转为ONNX模型很简单，使用torch.onnx.export()函数即可。
函数说明：

torch.onnx.export()
功能：
将.pth模型转为onnx文件导出

参数：

• model(torch.nn.Module): pth模型文件
• args (tuple of arguments): 模型的输入，模型的尺寸
• export_params (bool, default True):
  如果指定为True或者默认，参数也会被导出，如果要导出一个没训练过的就设置为False
• verbose (bool, default False): 导出轨迹的调试描述
• training (bool, default False) ：在训练模式下导出模型。目前，ONNX导出的模型只是为了做推断，通常不需要将其设置为True；
• input_names (list of strings, default empty list) ：onnx文件的输入名称, 可以随便取
• output_names (list of strings, default empty list) ：onnx文件的输出名称，可以随便取
• opset_version：默认为9
• dynamic_axes – {‘input’ : {0 : ‘batch_size’}, ‘output’ : {0 :
  ‘batch_size’}})

转换示例代码

import torch
import torch.nn as nn


# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__(