PyTorch Hub加载预训练模型|极客笔记

PyTorch Hub加载预训练模型

简介

PyTorch Hub是一个预训练模型库，让用户可以轻松地使用第三方发布的预训练模型。这些预训练模型包含了在大规模数据集上进行训练后的参数，可以用来进行迁移学习或者作为基础模型进行微调。

在本文中，我们将详细介绍如何使用PyTorch Hub加载预训练模型，并展示如何使用加载的模型进行推断。

加载预训练模型

要加载PyTorch Hub中的预训练模型，可以使用torch.hub.load()函数。该函数有三个参数：

repository: 模型所在的仓库名
model: 模型名
pretrained: 是否加载预训练参数（默认为True）

下面是一个示例，演示如何加载PyTorch Hub中的ResNet模型并进行推断：

import torch
model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True)
model.eval()

# Create a random input tensor
input_tensor = torch.rand(1, 3, 224, 224)

# Perform inference
output = model(input_tensor)
print(output)

运行以上代码，将加载ResNet-18模型并对随机输入张量进行推断。输出张量output将包含模型的推断结果。

PyTorch Hub中的流行模型

PyTorch Hub中包含了许多流行的预训练模型，包括图像分类、物体检测、文本生成等。以下是一些常见的模型：

pytorch/vision: 包含图像分类、目标检测等视觉任务的模型
pytorch/audio: 包含音频处理任务的模型
pytorch/text: 包含文本生成和文本分类任务的模型

使用torch.hub.list()函数可以查看PyTorch Hub中的所有模型列表。

迁移学习

一种常见的用法是使用PyTorch Hub中的预训练模型进行迁移学习。通过加载一个在大规模数据集上训练好的模型，我们可以在自己的数据集上进行微调，从而获得更好的性能。

以下是一个简单的示例，演示如何使用PyTorch Hub中的预训练ResNet模型进行迁移学习：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True)

# Replace the last fully connected layer with a new one
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10)  # Assuming we have 10 classes

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Train the model
for epoch in range(num_epochs):
    # Training loop
    for inputs, labels in dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

# Save the model
torch.save(model.state_dict(), 'model.pth')

在上面的示例中，我们加载了预训练的ResNet-18模型，并修改了最后一层全连接层以适应我们自己的数据集。然后我们定义了损失函数和优化器，并进行了训练。最后，我们保存了微调后的模型参数。