Pytorch 如何加载预训练的PyTorch模型

在本文中，我们将介绍如何使用PyTorch加载预训练的模型。预训练模型是指在大型数据集上进行训练并保存的模型，它们通常能够提供出色的性能，并且可以直接用于特定任务或进行微调。

1. 加载预训练模型

PyTorch提供了一个torchvision.models模块，其中包含了许多常用的预训练模型，如ResNet、VGG等。我们可以使用这些模型来加载预训练权重。

下面是一个加载预训练ResNet模型并使用的示例：

import torch
import torchvision.models as models

# 加载ResNet模型
resnet = models.resnet50(pretrained=True)

# 使用预训练模型进行推理
input_data = torch.randn(1, 3, 224, 224)
output = resnet(input_data)

在上面的代码中，首先我们导入了torch和torchvision.models模块。然后我们使用models.resnet50(pretrained=True)来加载预训练的ResNet-50模型。pretrained=True表示使用预训练的权重。接下来，我们定义了一个输入数据input_data，并将其作为参数传递给resnet模型。最后，我们通过调用resnet(input_data)来进行推理。

2. 加载自定义的预训练模型

除了加载torchvision提供的预训练模型外，我们还可以加载并使用自定义的预训练模型。下面是一个加载自定义预训练模型并使用的示例：

import torch
from torchvision import models

# 定义自定义模型的类
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.fc1 = torch.nn.Linear(64*56*56, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = x.view(-1, 64*56*56)
        x = self.fc1(x)
        return x

# 加载自定义模型
model = MyModel()
model.load_state_dict(torch.load('my_model.pth'))
model.eval()

# 使用预训练模型进行推理
input_data = torch.randn(1, 3, 224, 224)
output = model(input_data)

在上面的代码中，我们首先定义了一个自定义模型MyModel，该模型包含卷积层和全连接层。然后，我们加载了已保存的自定义模型权重，并通过调用model.load_state_dict(torch.load(‘my_model.pth’))将权重加载到模型中。此外，我们还通过调用model.eval()将模型设置为评估模式，以确保不进行训练。

最后，我们定义了输入数据input_data，并将其作为参数传递给model模型。通过调用model(input_data)，我们可以使用自定义的预训练模型进行推理。

总结

本文介绍了如何使用PyTorch加载预训练的模型。通过使用torchvision.models模块加载预训练模型，我们可以方便地使用常用的模型，如ResNet、VGG等。此外，我们还可以加载并使用自定义的预训练模型，通过调用model.load_state_dict(torch.load(‘my_model.pth’))将权重加载到自定义模型中。加载预训练模型后，我们可以使用其进行推理或进行微调，以满足特定的任务需求。

希望本文对你理解如何加载预训练的PyTorch模型有所帮助！