Torch中高效利用预训练权重：从加载到微调

引言

在深度学习中，预训练模型是一种强大的工具，它们通过在大规模数据集（如ImageNet）上进行训练，已经学习到了丰富的特征表示。利用这些预训练模型，我们可以在较小的数据集上快速训练出高性能的模型，或者通过微调（fine-tuning）来适应新的任务。PyTorch作为一个流行的深度学习框架，提供了方便的接口来加载和使用预训练模型。

1. 加载预训练模型

PyTorch的torchvision库包含了许多流行的预训练模型，如ResNet、VGG、AlexNet等。以下是一个加载ResNet18预训练权重的示例：

import torchvision.models as models
# 加载预训练模型
model = models.resnet18(pretrained=True)
# 将模型设置为评估模式（如果你不打算训练它）
model.eval()
# 如果你想看到模型的参数和结构
print(model)

在上面的代码中，pretrained=True参数告诉PyTorch自动下载并加载预训练的权重。如果你只想获取模型结构而不加载预训练权重，可以设置为pretrained=False。

2. 修改模型以适应新任务

预训练模型通常用于图像分类任务，其输出层对应于训练时使用的类别数。如果你的任务类别数不同，你需要修改模型的最后几层。以下是一个修改ResNet18输出层以匹配新类别数的例子：

num_ftrs = model.fc.in_features  # 获取原始全连接层的输入特征数
model.fc = torch.nn.Linear(num_ftrs, num_classes)  # 替换全连接层

其中num_classes是你的新任务中的类别数。

3. 微调模型

微调是调整预训练模型参数以适应新任务的过程。在微调时，通常会冻结模型的一部分层（通常是前面的层），只训练后面的层。这有助于保留模型学习到的通用特征，同时让模型学习新任务特有的特征。

# 冻结模型的大部分层
for param in model.parameters():
    param.requires_grad = False
# 只解冻最后的全连接层
model.fc.requires_grad = True
# 如果你使用的是优化器，确保只包含需要梯度的参数
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)
# 训练循环...

4. 训练和评估

一旦模型准备好，你就可以使用自己的数据集进行训练和评估了。这通常涉及到前向传播、计算损失、反向传播和更新权重等步骤。

# 假设你有一个数据加载器data_loader
for inputs, labels in data_loader:
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = loss_fn(outputs, labels)
    loss.backward()
    optimizer.step()
# 评估模型...

5. 注意事项

• 学习率：微调时，使用比从头训练时更低的学习率通常是有益的。
• 冻结层：根据任务的不同，你可能需要调整冻结的层数。对于非常不同的任务，可能需要解冻更多的层。
• 数据预处理：确保你的数据预处理方式与预训练模型训练时使用的方式一致。

结论

通过使用预训练模型并在自己的数据集上进行微调，你可以显著加速训练过程，并可能获得更好的模型性能。PyTorch提供了强大的工具和库来支持这一过程，使得即使是非专业用户也能轻松上手。希望本文能帮助你更好地理解和使用PyTorch中的预训练模型。