基于pytorch框架的lenet模型并导出ONNX格式

import torch.nn as nn
import torch.nn.functional as F
import torch
 
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, 5)
        self.conv3 = nn.Conv2d(20, 40, 3)
        self.mp = nn.MaxPool2d(2)
        self.mp1 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(25000, 512)
        self.fc2 = nn.Linear(512, 10)
 
    def forward(self, x):
        in_size = x.size(0)
        x = F.relu(self.mp(self.conv1(x)))
        x = F.relu(self.mp(self.conv2(x)))
        x = F.relu(self.mp1(self.conv3(x)))
        x = x.view(in_size, -1)
        x = self.fc1(x)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)
 
 
model = Net().to(device)

# torch.save(model, './model_para.pth')
torch.save(model, './model_para.pth')

torch_model = torch.load("./model_para.pth") # pytorch模型加载
batch_size = 1  #批处理大小
input_shape = (1,224,224)   #输入数据
 
# set the model to inference mode
torch_model.eval()
 
x = torch.randn(batch_size,*input_shape)		# 生成张量
export_onnx_file = "test.onnx"		        	# 目的ONNX文件名
torch.onnx.export(torch_model,
                    x,
                    export_onnx_file,
                    opset_version=10,
                    do_constant_folding=True,	# 是否执行常量折叠优化
                    input_names=["input"],		# 输入名
                    output_names=["output"],	# 输出名
                    dynamic_axes={"input":{0:"batch_size"},		# 批处理变量
                                    "output":{0:"batch_size"}})


#torch.save(model.state_dict(), './model_para.pth')

导出模型文件.pth格式和ONNX格式

参考资料：

pytorch可视化工具netron - 江南烟雨尘 - 博客园

---

结束！