最近pytorch转onnx模型时自己从头训练了mlp模型,然后使用不同的方式保存成pth,最后使用torch.onnx成功转出onnx模型,注意本例不考虑dynamic shape的情形,后面有机会再单独写。
import os
import torch
from torch import nn
from torchvision.datasets import CIFAR10
from torch.utils.data import DataLoader
from torchvision import transforms
import torch.onnx
class MLP(nn.Module):
'''
Multilayer Perceptron.
'''
def __init__(self):
super().__init__()
self.layers = nn.Sequential(
nn.Flatten(),
nn.Linear(32 * 32 * 3, 64),
nn.ReLU(),
nn.Linear(64, 10),
)
def forward(self, x):
'''Forward pass'''
return self.layers(x)
def test_model_load():
modle_origin = torch.load('mlp_origin.pth')
modle = torch.load('mlp_whole.pth')
print(modle)
weights = torch.load('mlp_weights.pth')
print(weights)
return None
def test_onnx_export():
# modle = torch.load('mlp_whole.pth') #直接load整个模型,前提是保存的pth要是包含网络
modle = MLP() #创建网络
modle.load_state_dict(torch.load('mlp_weights.pth')) #load网络对应节点的权重
modle.eval() #转换成inference模式
# Input to the model
x = torch.randn(1, 3, 32, 32, requires_grad=True)
torch_out = modle(x)
# Export the model
torch.onnx.export(modle, # model being run
x, # model input (or a tuple for multiple inputs)
"mlp_test1.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=12, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names=['input'], # the model's input names
output_names=['output'], # the model's output names
# dynamic_axes={'input': {0: 'batch_size'}, # variable length axes
# 'output': {0: 'batch_size'}}
)
if __name__ == '__main__':
# # test_model_load()
# test_onnx_export()
# exit
# Set fixed random number seed
torch.manual_seed(42)
# Prepare CIFAR-10 dataset
dataset = CIFAR10(os.getcwd(), download=False, transform=transforms.ToTensor())
trainloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True, num_workers=1)
# Initialize the MLP
mlp = MLP()
torch.save(mlp, 'mlp_origin.pth')
# Define the loss function and optimizer
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(mlp.parameters(), lr=1e-4)
# Run the training loop
for epoch in range(0, 2): # 5 epochs at maximum
# Print epoch
print(f'Starting epoch {epoch + 1}')
# Set current loss value
current_loss = 0.0
# Iterate over the DataLoader for training data
for i, data in enumerate(trainloader, 0):
# Get inputs
inputs, targets = data
# Zero the gradients
optimizer.zero_grad()
# Perform forward pass
outputs = mlp(inputs)
# Compute loss
loss = loss_function(outputs, targets)
# Perform backward pass
loss.backward()
# Perform optimization
optimizer.step()
# Print statistics
current_loss += loss.item()
if i % 500 == 499:
print('Loss after mini-batch %5d: %.3f' %
(i + 1, current_loss / 500))
current_loss = 0.0
# Process is complete.
torch.save(mlp.state_dict(), 'mlp_weights.pth')
torch.save(mlp, 'mlp_whole.pth')
print('Training process has finished.')
mlp部分参考自:https://github.com/christianversloot/machine-learning-articles/blob/main/creating-a-multilayer-perceptron-with-pytorch-and-lightning.md