开始配置网络,由于ImageNet是识别1000个物体,我们的狗的分类一共只有120,所以需要对模型的最后一层全连接层进行微调,将输出从1000改为120
model_ft = models.resnet50(pretrained=True) # 这里自动下载官方的预训练模型,并且
# 将所有的参数层进行冻结
for param in model_ft.parameters():
param.requires_grad = False
# 这里打印下全连接层的信息
print(model_ft.fc)
num_fc_ftr = model_ft.fc.in_features #获取到fc层的输入
model_ft.fc = nn.Linear(num_fc_ftr, len(breeds)) # 定义一个新的FC层
model_ft=model_ft.to(DEVICE)# 放到设备中
print(model_ft) # 最后再打印一下新的模型
参考:https://pytorchbook.cn/chapter4/4.1-fine-tuning/#