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 - 江南烟雨尘 - 博客园