PyTorch 实战：加载模型权重与 ONNX 推理部署

使用 ONNX Runtime 进行预测

导出成功后，我们可以用 onnxruntime 来加载模型并跑通一次完整的预测流程。这里以 FashionMNIST 数据集为例，取第一张图进行测试。

import onnxruntime
from torchvision import datasets
from torchvision.transforms import ToTensor

test_data = datasets.FashionMNIST(
    root="data", train=False, download=True, transform=ToTensor()
)
classes = [
    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]

x, y = test_data[0][0], test_data[0][1]
session = onnxruntime.InferenceSession(onnx_model, None)

input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
result = session.run([output_name], {input_name: x.numpy()})

predicted = classes[result[0][0].argmax(0)]
actual = classes[y]
print(f'Predicted: "{predicted}", Actual: "{actual}"')

输出结果如下：

Predicted: "Ankle boot", Actual: "Ankle boot"

可以看到，模型成功识别出了短靴。通过这种方式，我们实现了在不同环境下复用训练好的模型。

完整代码参考

为了方便大家直接运行，这里汇总了上述所有步骤的代码。

import torch
import onnxruntime
from torch import nn
import torch.onnx as onnx
from torchvision import datasets
from torchvision.transforms import ToTensor

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
            nn.ReLU()
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

# 加载模型
model = NeuralNetwork()
model.load_state_dict(torch.load('data/model.pth'))
model.eval()

# 导出 ONNX
input_image = torch.zeros((1, 28, 28))
onnx.export(model, input_image, 'data/model.onnx')

# 推理
from torchvision import datasets
from torchvision.transforms import ToTensor

test_data = datasets.FashionMNIST(
    root="data", train=False, download=True, transform=ToTensor()
)
classes = [
    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]

x, y = test_data[0][0], test_data[0][1]
session = onnxruntime.InferenceSession('data/model.onnx', None)
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
result = session.run([output_name], {input_name: x.numpy()})

predicted = classes[result[0][0].argmax(0)]
actual = classes[y]
print(f'Predicted: "{predicted}", Actual: "{actual}"')

总结

本节完成了模型从加载到跨平台推理的闭环。我们不仅掌握了 load_state_dict 和 eval() 的使用细节，还学会了利用 ONNX 解决部署兼容性问题。结合之前的张量操作、数据预处理和自动微分知识，现在你已经具备了构建基本图像分类模型的能力。