本文共 1512 字，大约阅读时间需要 5 分钟。

说明

本文代码来自开源项目生 Wilkinson 是对代码进行了大量注释以便理解各个函数的原理和用途。

配置环境

使用环境：Python 3.8

代码致谢

感谢生 Wilkinson 的开源项目，我对这段代码进行了注释。

图像增广部分

由于这节相对复杂，代码注释量较多

预处理步骤

为了获得好的训练效果，我们对输入图像进行了以下处理

模型配置

使用了预训练的ResNet-18模型

pretrained_net = models.resnet18(pretrained=True)print(pretrained_net.fc)print(pretrained_net)

权重调整

由于输出类别数为2，我们需要修改前好的全连接层

pretrained_net.fc = nn.Linear(512, 2)print(pretrained_net.fc)print(pretrained_net)

微调学习率设置

fc层的参数初始是随机的需要较大的学习率来训练

lr = 0.01feature_params = list(map(id, pretrained_net.fc.parameters()))# 筛选出不属于fc层的参数feature_params = filter(lambda p: id(p) not in feature_params, pretrained_net.parameters())

优化器设置

使用SGD优化器，对特征参数和全连接参数分别设置学习率

optimizer = optim.SGD([    {'params': feature_params},    {'params': pretrained_net.fc.parameters(), 'lr': lr * 10}], lr=lr, weight_decay=0.001)

训练函数

这是一个完整的训练函数，可以复用于其他模型

def train_fine_tuning(net, optimizer, batch_size=128, num_epochs=5):    train_iter = DataLoader(...)    test_iter = DataLoader(...)    criterion = loss = CrossEntropyLoss()    d2l.train(train_iter, test_iter, net, loss, optimizer, device, num_epochs)

测试代码

为了对比，我们定义了一个完全从头开始训练的模型

scratch_net = models.resnet18(pretrained=False, num_classes=2)optimizer = optim.SGD(scratch_net.parameters(), lr=0.1, weight_decay=0.001)train_fine_tuning(scratch_net, optimizer)

这段代码主要展示了如何在PyTorch中使用预训练的ResNet进行迁移学习，同时对模型进行微调。