torch.onnx.export示例

以下是一个使用`torch.onnx.export`函数的简单示例：

import torch
import torch.onnx as onnx

# 定义一个简单的模型
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 6, 5)
        self.pool = torch.nn.MaxPool2d(2, 2)
        self.conv2 = torch.nn.Conv2d(6, 16, 5)
        self.fc1 = torch.nn.Linear(16 * 5 * 5, 120)
        self.fc2 = torch.nn.Linear(120, 84)
        self.fc3 = torch.nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(torch.nn.functional.relu(self.conv1(x)))
        x = self.pool(torch.nn.functional.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 创建一个实例
net = Net()

# 加载一个预训练模型
model_path = 'model.pth'
net.load_state_dict(torch.load(model_path))

# 将模型转换为ONNX格式
input_names = ['input']
output_names = ['output']
dummy_input = torch.randn(1, 1, 28, 28)
onnx_path = 'model.onnx'
onnx.export(net, dummy_input, onnx_path, input_names=input_names, output_names=output_names)

print('Model converted to ONNX format.')

该示例定义了一个简单的卷积神经网络模型，加载了预训练的权重，并使用`torch.onnx.export`函数将模型转换为ONNX格式。需要注意的是，模型转换为ONNX格式时需要指定输入和输出的名称，以便后续在使用ONNX模型时能够正确地处理输入和输出。