day8-模型微调

本是应用于迁移学习的方法,使用预训练的模型,但是输出层自己定义

定义并加载模型

model_conv = torchvision.models.resnet18(pretrained=True)
for param in model_conv.parameters():
    param.requires_grad = False

# Parameters of newly constructed modules have requires_grad=True by default
num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 2)

model_conv = model_conv.to(device)

剩下就是读取数据正常训练了

如果不冻结前面层的参数,建议前面层lr较小,输出层lr较大

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

你可能感兴趣的:(day8-模型微调)