PyTorch 模型加载

如果你使用 PyTorch 框架训练了模型，那么在使用该模型时，需要将其加载到内存中。本文将介绍如何加载 PyTorch 模型。

准备工作

假设你已经完成了 PyTorch 模型的训练，并将模型以 .pt 或 .pth 格式保存在磁盘上。此外，你还需要安装 PyTorch 库。你可以在命令行中运行以下命令来安装 PyTorch：

pip install torch

加载模型

你可以使用 torch.load 函数将 PyTorch 模型加载到内存中。该函数的基本语法如下：

model = torch.load(PATH)

其中，PATH 是模型的文件路径。加载模型后，你就可以将其应用到新的数据上进行预测。

示例

以下示例演示如何使用 torch.load 函数加载 MNIST 数据集上训练得到的简单卷积神经网络模型：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.mp = nn.MaxPool2d(2)
        self.fc = nn.Linear(320, 10)

    def forward(self, x):
        in_size = x.size(0)
        x = self.mp(self.conv1(x))
        x = self.mp(self.conv2(x))
        x = x.view(in_size, -1)
        x = self.fc(x)
        return x

# 加载模型
model = Net()
model.load_state_dict(torch.load("model.pt"))
model.eval()

# 加载测试数据集
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor()))

# 测试模型
test_loss = 0
correct = 0
with torch.no_grad():
    for data, target in test_loader:
        output = model(data)
        test_loss += F.nll_loss(output, target, reduction='sum').item()
        pred = output.argmax(dim=1, keepdim=True)
        correct += pred.eq(target.view_as(pred)).sum().item()

test_loss /= len(test_loader.dataset)
print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
    test_loss, correct, len(test_loader.dataset),
    100. * correct / len(test_loader.dataset)))

我们假设该程序文件名为 test.py，模型文件名为 model.pt。你可以通过以下命令运行该程序：

python test.py

当程序运行结束后，你将看到以下输出结果：

Test set: Average loss: 0.0347, Accuracy: 9885/10000 (99%)

说明加载的模型在测试集上取得了 99% 的准确率。

小结

本文介绍了如何使用 PyTorch 加载模型。你可以参考本文中的示例代码在自己的项目中使用 PyTorch 模型。