24,PyTorch 预训练模型的加载与使用

乔丹搞IT

于 2025-06-29 16:33:55 发布

阅读量198

点赞数 14

CC 4.0 BY-SA版权

CC BY-NC-SA 3.0

本文链接：https://blog.csdn.net/jimn2000/article/details/148999709

PyTorch 专栏收录该内容

23 篇文章

订阅专栏

在这里插入图片描述

24. PyTorch 预训练模型的加载与使用

在深度学习中，预训练模型是迁移学习的核心，它们为各种任务提供了强大的特征提取能力。PyTorch 提供了丰富的预训练模型库，这些模型大多在大规模数据集（如 ImageNet）上进行了训练，能够直接用于多种任务，或者通过微调来适应特定的应用场景。本文将详细介绍如何在 PyTorch 中加载和使用这些预训练模型，以及如何根据具体需求进行调整和优化。

24.1 PyTorch 预训练模型的来源

PyTorch 的预训练模型主要来源于 torchvision.models 模块。这些模型包括了经典的卷积神经网络（CNN）架构，如 ResNet、VGG、AlexNet 等，以及一些更现代的架构，如 EfficientNet 和 Vision Transformer（ViT）。这些模型在 ImageNet 数据集上进行了预训练，能够提取通用的图像特征，适用于各种视觉任务。

24.2 加载预训练模型

加载预训练模型非常简单，只需从 torchvision.models 模块中选择所需的模型，并设置 pretrained=True 即可。以下是一个加载 ResNet50 模型的示例：

import torchvision.models as models

# 加载预训练的 ResNet50 模型
model = models.resnet50(pretrained=True)

加载模型后，可以直接查看其结构，了解各层的参数和功能。例如：

print(model)

这将输出模型的详细结构，包括每一层的名称、类型和参数数量。

24.3 使用预训练模型进行特征提取

预训练模型的一个重要用途是作为特征提取器。通过冻结模型的权重，可以直接使用其前几层提取的特征，而无需对模型进行微调。这在处理小规模数据集或需要快速原型开发时非常有用。

以下是一个使用预训练模型进行特征提取的示例：

import torch

# 冻结模型的所有参数
for param in model.parameters():
    param.requires_grad = False

# 定义一个特征提取函数
def extract_features(model, input_tensor):
    model.eval()  # 设置为评估模式
    with torch.no_grad():  # 禁用梯度计算
        features = model(input_tensor)
    return features

# 假设 input_tensor 是输入图像的张量
input_tensor = torch.randn(1, 3, 224, 224)  # 示例输入
features = extract_features(model, input_tensor)

print(features.shape)  # 输出特征的形状

在这个例子中，extract_features 函数将输入图像通过预训练模型，提取出特征向量。这些特征可以直接用于分类、聚类或其他任务。

24.4 微调预训练模型

虽然预训练模型在通用数据集上表现良好，但在特定任务上可能需要进一步优化。微调是迁移学习中常用的技术，通过在目标数据集上对预训练模型进行少量的训练，可以显著提高模型的性能。

微调通常包括以下步骤：

冻结部分层：保留预训练模型的大部分层，冻结其权重，以保留通用特征。
替换最后的分类层：根据目标任务的类别数量，替换预训练模型的最后分类层。
训练模型：在目标数据集上对模型进行训练，优化其性能。

以下是一个完整的微调示例：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision import transforms
import torchvision.models as models

# 加载预训练模型
model = models.resnet50(pretrained=True)

# 冻结部分层
for param in model.parameters():
    param.requires_grad = False

# 替换最后的分类层
num_classes = 10
model.fc = nn.Linear(model.fc.in_features, num_classes)

# 定义数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据集
train_dataset = CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)

# 微调模型
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')

在这个例子中，我们使用了 CIFAR-10 数据集对 ResNet50 模型进行了微调。通过冻结大部分层并替换最后的分类层，模型能够快速适应新的任务。

24.5 选择合适的预训练模型

选择合适的预训练模型对于迁移学习的成功至关重要。不同的模型在不同的任务上表现可能不同，因此需要根据具体任务选择合适的模型。以下是一些常见的预训练模型及其适用场景：

ResNet：适用于图像分类任务，具有良好的泛化能力和较高的准确率。
VGG：适用于图像分类和目标检测任务，结构简单，易于实现。
AlexNet：是早期的经典模型，适用于图像分类任务，但计算量较大。
EfficientNet：在计算效率和准确率之间取得了良好的平衡，适用于资源受限的设备。
Vision Transformer (ViT)：适用于图像分类和视觉任务，具有强大的特征提取能力，但对数据量要求较高。

24.6 实际应用中的注意事项

在实际应用中，使用预训练模型时需要注意以下几点：

数据预处理：确保输入数据的格式和预训练模型的输入要求一致，包括图像大小、归一化参数等。
学习率调整：在微调过程中，学习率的选择非常重要。通常，对于冻结的层，学习率可以设置得较低；而对于新添加的层，学习率可以设置得较高。
数据增强：通过使用数据增强技术（如随机裁剪、翻转、旋转等），可以进一步提高模型的泛化能力，使其更好地适应目标任务。
硬件资源：预训练模型通常较大，需要足够的内存和计算资源。如果资源有限，可以考虑使用轻量级模型，如 MobileNet 或 EfficientNet。

24.7 总结

PyTorch 提供了丰富的预训练模型，这些模型为迁移学习提供了强大的支持。通过加载预训练模型、冻结部分层、替换分类层和微调模型，可以在各种任务中快速构建和优化模型。选择合适的预训练模型并合理调整训练策略，可以显著提高模型的性能和泛化能力。希望本文能够帮助你更好地理解和使用 PyTorch 中的预训练模型，提升你的深度学习项目的效果。
更多技术文章见公众号: 大城市小农民