0.前言

        写这篇文章的时候其实是当时在写模型量化的时候，突然想到在使用qnn的时候好像需要中间有一部转成onnx。很多小伙伴对这个onnx又很陌生，以及pth转成onnx的时候各种导出参数含义又很模糊。所以就先写一下这个东西吧。然后查看这个onnx又有一个非常好用的可视化界面netron，可以清晰的看到模型的各种结构和参数！！网页地址如下 Netron

1.为啥要转onnx

        这时候有小伙伴说，我们讲的不是netron吗，怎么说到onnx了。因为netron对pth权重文件非常不友好啊，请看图。。。。。非常的鸡肋，什么都看不见。

        其实pth转成onnx这个非常的简单，也就是一个转化代码的事情torch.onnx.export（），封装的非常完美。

2.如何导出onnx

这边举一个例子啊，看代码

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(64 * 32 * 32, 10)
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x
model = SimpleModel()
model.eval()  

dummy_input = torch.randn(1, 3, 32, 32)

# 导出ONNX！！！
torch.onnx.export(
    model,                              # 要导出的模型
    dummy_input,                        # 示例输入
    "./model2.onnx",                    # 输出文件名
    export_params=True,                 # 导出训练好的参数
    opset_version=11,                   # ONNX算子集版本
    do_constant_folding=True,           # 是否执行常量折叠优化
    input_names=['input'],              # 输入名称
    output_names=['output'],            # 输出名称
    dynamic_axes={                      # 动态轴（可变批量大小）
        'input': {0: 'batch_size'},
        'haodong': {0: 'batch_size'}
    },
)
print("模型已成功导出为")

        这样就可以很轻松的导出onnx，我导出的路径就在当前文件夹下。然后把这个onnx文件直接拖入到页面中就可以啦，然后你就可以看到这个东西。

3.netron的使用

其实到这里netron的核心操作基本就结束了。还有一些布局操作可以自己摸索一下。

4.导出onnx文件的说明

        不知道各位有没有发现，我在导出onnx这个步骤的时候，输入的参数还是比较多的，因为它对应了不同模型的导出。

torch.onnx.export(
    model,                              
    dummy_input,                        
    "./model2.onnx",                   
    export_params=True,                
    opset_version=11,                  
    do_constant_folding=True,           
    input_names=['input'],              # 输入名称
    output_names=['output'],            # 输出名称
    dynamic_axes={                      # 动态轴（可变批量大小）
        'input': {0: 'batch_size'},
        'haodong': {0: 'batch_size'}
    },
)

        主要就讲解这个名字和动态轴的含义。这个input_names和output_names主要对应你的输入输出，她是可以多输入和多输出的！假如你的模型是一个输入，四个输出，你是可以随意给他们定义名称的，比如我下面随便搞的一个模型，输入的节点我定义为'images'，输出四个节点定义为'classification', 'regression', 'combined', 'features'。

import torch
import torch.nn as nn

class ComplexIOModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3, padding=1)  
        self.fc1 = nn.Linear(64 * 32 * 32, 10)  
        self.fc2 = nn.Linear(64 * 32 * 32, 5)
        
    def forward(self, image):    
        img_feat = self.conv(image)  
        img_feat = img_feat.view(img_feat.size(0), -1)
        class_scores = self.fc1(img_feat)  
        reg_values = self.fc2(img_feat)   
        combined = torch.cat([class_scores, reg_values], dim=1) 
        
        return class_scores, reg_values, combined, img_feat

model = ComplexIOModel()
model.eval()


dummy_image = torch.randn(1, 3, 32, 32)

torch.onnx.export(
    model, 
    dummy_image, 
    "complex_model.onnx",
    export_params=True,
    opset_version=11,
    input_names=['image'],                 
    #一个输入！！！
    output_names=['classification', 'regression', 'combined', 'features'],
    #四个输出！！！
    dynamic_axes={
        'image': {0: 'batch_size'},
        'classification': {0: 'batch_size'},
        'regression': {0: 'batch_size'},
        'combined': {0: 'batch_size'},
        'features': {0: 'batch_size'}
    }
)

print(" 模型导出成功！")

        细心的小伙伴会发现，那个动态轴也跟着在变化。那是什么意思呢。对于pytorch框架训练的模型，它的维度遵循[batch, c, h, w],你可以想象为一个很多的立方体。然后你可以对这个输出进行索引比如batch就是第一个维度[0], 就是关于立方体的数量。c维度就是[1]，就是关于立方体的深度。。。以此类推。说到这其实很多小伙伴对这个张量的运算其实有疑惑的，我后面也会出一期关于张量运输的说明，这个东西还是挺抽象的，四个维度其实还是有含义，可以想象出来，维度很高的时候就难以想象，只能遵循一定的运算法则来运算，到时候我来讲transformer的时候着重讲一下张量是咋运算的。

        回到主题，这时候我们就知道，这个0就是对应的batch，对于一个固定的网络来说，batch是变化的，但是onnx导出的时候如果你不说，他就会对这个东西进行固定，比如说，我们举例子的时候，dummy_image = torch.randn(1, 3, 32, 32)。如果不指定动态轴，你下一次输出(2, 3, 32, 32)，他就会报错。这样可以吗？肯定不可以啊，模型推理的时候也是一个batch一个batch的推理啊。当然，其他维度你也可以根据模型的需要进行设置。

5.题外话：

        原本写这个博客就是图一乐，大家理性观看，本人也是学生，水平有限，有些写的不对的地方可以相互交流。再加上每次写这个博客都是晚上，干了一天活，一肚子牢骚。必须狠狠地在博客里面乱写一通啊。什么时候当上老板，不用上班学习了，就把笔记本狠狠砸烂，然后一天疯狂睡10个小时，让手底下的实习生狠狠的加班这一块。