如何在 PyTorch 中冻结模型权重以进行迁移学习：分步教程

在PyTorch中，对预训练模型进行迁移学习时，通常需要冻结一部分权重以防止它们在训练新任务时发生意外改变。以下是分步教程，演示如何在一个预训练的ResNet模型上进行权重冻结：

1. **导入所需库**：

   import torch
   from torch import nn
   from torchvision.models import resnet18

2. **加载预训练模型**：

   model = resnet18(pretrained=True)

`pretrained=True`会自动下载并加载ImageNet数据集训练的预训练权重。

3. **查看模型结构**：

   print(model)

确定哪些层是可训练（`requires_grad=True`）的，哪些是不可训练（`requires_grad=False`）的。

4. **冻结所有预训练层**：

   for param in model.parameters():
       param.requires_grad = False

这将设置所有预训练层的权重为不可训练，只允许新的自定义添加的层学习。

5. **选择要在训练阶段更新的自定义层**：

   custom_layers = list(model.children())[-num_custom_layers:]  # 根据你的需求指定层数
   for layer in custom_layers:
       for param in layer.parameters():
           param.requires_grad = True

6. **定义损失函数、优化器和学习率调度器**：

   criterion = nn.CrossEntropyLoss()
   optimizer = torch.optim.Adam(custom_layers.parameters(), lr=learning_rate)
   scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=steps_per_epoch, gamma=gamma)

7. **开始训练**：

   for epoch in range(num_epochs):
       for inputs, labels in dataloader:
           inputs, labels = inputs.to(device), labels.to(device)
           outputs = model(inputs)
           loss = criterion(outputs, labels)

           optimizer.zero_grad()  # 清空梯度
           loss.backward()  # 计算梯度
           optimizer.step()  # 更新参数

           # 学习率调度
           scheduler.step()

   # 结束训练后，如果想要解冻全部参数进行微调，可以这样做：
   for param in model.parameters():
       param.requires_grad = True

8. **评估和微调**（如果你需要）：
在训练完成后，你可以解冻整个模型进行进一步的微调，或仅对特定层进行微调。