案例一:猫狗分类

参考:猫狗分类迁移学习案例详解
代码位置:E:\项目例程\猫狗分类\迁移学习\猫狗_resnet18_2 \猫狗分类_迁移学习可视化

参数微调

该方法使用预训练的参数来初始化我们的网络模型,修改全连接层后再训练所有层。

# 加载预训练模型
model_ft = models.resnet18(pretrained=True)

# 获取resnet18的全连接层的输入特征数
num_ftrs = model_ft.fc.in_features

# 调整全连接层的输出特征数为2
model_ft.fc = nn.Linear(num_ftrs, len(class_names))

# 将模型放到GPU/CPU
model_ft = model_ft.to(device)

# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 选择优化器
optimizer_ft = optim.SGD(model_ft.parameters(), lr=1e-3, momentum=0.9)

# 定义优化器器调整策略,每5轮后学习率下调0.1个乘法因子
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=5, gamma=0.1)

# 调用训练函数训练
model_ft = train_model(
    model_ft, 
    criterion, 
    optimizer_ft, 
    exp_lr_scheduler,
    num_epochs=10
)

特征提取的方法

该方法冻结除全连接层外的所有层的权重,修改全连接层后仅训练全连接层。

# 加载预训练模型
model_conv = models.resnet18(pretrained=True)

# 冻结除全连接层外的所有层, 使其梯度不会在反向传播中计算
for param in model_conv.parameters():
    param.requires_grad = False

# 获取resnet18的全连接层的输入特征数
num_ftrs = model_conv.fc.in_features

# 调整全连接层的输出特征数为2
model_conv.fc = nn.Linear(num_ftrs, 2)

# 将模型放到GPU/CPU
model_conv = model_conv.to(device)

# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 选择优化器, 只传全连接层的参数
optimizer_conv = optim.SGD(model_conv.fc.parameters(), lr=1e-3, momentum=0.9)

# 定义优化器器调整策略,每5轮后学习率下调0.1个乘法因子
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=5, gamma=0.1)

# 调用训练函数训练
model_conv = train_model(
    model_conv,
    criterion,
    optimizer_conv,
    exp_lr_scheduler,
    num_epochs=10
)

总结

对比发现,使用迁移学习起点更高,收敛更快。其中特征提取总用时更短,准确率更高。这个结果是预料之中的,在迁移方法的选择中,我们知道对于数据集较小,且与原始数据集相似度高时,选择特征提取的方法会更好。

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐