pytorch模型转为onnx

一 代码实现

import torch
import torchvision


def main():

    # 模型结构加载
    model = torchvision.models.resnet50()
    # 修改层结构：输入图像为6通道，10个类
    # model.conv1 = nn.Conv2d(6, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
    # model.fc =  nn.Linear(in_features=2048, out_features=10, bias=True)

    # 加载本地训练权重/
    # pre_trained = 'model/resnet50.pth' 
    # pre_trained_dict = torch.load(pre_trained, map_location='cpu')
    # model.load_state_dict(pre_trained_dict)

    # 保证 BN 层能够用 全部训练数据 的均值和方差，即测试过程中要保证 BN 层的均值和方差不变
    model.eval()

    # 当前计算不需要反向传播
    with torch.no_grad():
        batch_size = 8

        # 按照标准正态分布生成模型输入数据
        dummy_input = torch.randn(batch_size,3,224,224)

        # 为了确保在推理过程中有不同批次的数据输入，这里设置动态尺寸 batch_size
        # 输入和输出的第 0 维度可以是动态，即 batch_size
        dynamic_axes = {'image': {0: 'batch_size'},
                        'class': {0: 'batch_size'}} 
        
        # model : 模型
        # dummy_input: 输入
        # output_file_path : ONNX保存路径
        # input_names: 输入节点名称
        # output_names: 输出节点名称

        output_file_path = "E:\\分类\\resnet50.onnx"

        torch.onnx.export(model, dummy_input, output_file_path,
                            verbose=True, input_names=['image'], output_names=['class'],
                            dynamic_axes=dynamic_axes)
    

if __name__ == '__main__':
    main()

二 结构分析

在网址https://netron.app/打开 resnet50.onnx