pytorch实现从本地加载 .pth 格式模型

在PyTorch中，我们可以使用.pth格式保存模型的权重和参数。在本文中，我们将详细讲解如何从本地加载.pth格式的模型。我们将使用两个示例来说明如何完成这些步骤。

示例1：加载全连接神经网络模型

以下是加载全连接神经网络模型的步骤：

import torch
import torch.nn as nn

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = Net()

# 加载模型权重
model.load_state_dict(torch.load('model.pth'))

# 使用模型进行预测
with torch.no_grad():
    inputs = torch.randn(1, 784)
    outputs = model(inputs)
    print(outputs)

在上述代码中，我们首先定义了一个简单的全连接神经网络Net，它含有一个输入层、一个隐藏层和一个输出层。然后，我们创建了一个模型实例model。我们使用model.load_state_dict()加载模型的权重，并使用with torch.no_grad()来禁用梯度计算，因为我们不需要计算梯度或更新权重。在使用模型进行预测时，我们使用torch.randn()函数生成一个随机输入，并将其传递给模型。

示例2：加载卷积神经网络模型

以下是加载卷积神经网络模型的步骤：

import torch
import torch.nn as nn

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=5, padding=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=5, padding=2)
        self.fc1 = nn.Linear(7 * 7 * 64, 1024)
        self.fc2 = nn.Linear(1024, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = x.view(-1, 7 * 7 * 64)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = Net()

# 加载模型权重
model.load_state_dict(torch.load('model.pth'))

# 使用模型进行预测
with torch.no_grad():
    inputs = torch.randn(1, 1, 28, 28)
    outputs = model(inputs)
    print(outputs)

在上述代码中，我们首先定义了一个简单的卷积神经网络Net，它含有两个卷积层、两个池化层和一个全连接层。然后，我们创建了一个模型实例model。我们使用model.load_state_dict()加载模型的权重，并使用with torch.no_grad()来禁用梯度计算，因为我们不需要计算梯度或更新权重。在使用模型进行预测时，我们使用torch.randn()函数生成一个随机输入，并将其传递给模型。

结论

在本文中，我们详细讲解了如何从本地加载.pth格式的模型。我们使用了两个示例来说明如何完成这些步骤。如果您按照这些步骤操作，您应该能够成功加载模型并使用它们进行预测。