我一直在关注Pythorch官方文档(http://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html)中的蚂蚁和蜜蜂转移学习教程。我正试图通过改变最后一层来预测两个类中的一个来优化VGG19模型。我可以使用以下代码修改最后一个fc层。在
但是我在执行train_model函数时遇到了一个错误。错误是“在/opt/conda/conda bld/pytorch_1513368888240/work/torch/lib/THC/generic处的大小不匹配/THC传感器布拉斯铜:243”. 知道是什么问题吗?在
model_conv = torchvision.models.vgg19(pretrained=True)
for param in model_conv.parameters():
param.requires_grad = False
model_conv = nn.Sequential(*list(model_conv.classifier.children())[:-1] +
[nn.Linear(in_features=4096, out_features=2)])
if use_gpu:
model_conv = model_conv.cuda()
criterion = nn.CrossEntropyLoss()
optimizer_conv = optim.SGD(model_conv._modules['6'].parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)
model_conv = train_model(model_conv, criterion, optimizer_conv, exp_lr_scheduler, num_epochs=25)
当您定义模型时,您只考虑
classifier
,它只包含在网络的完全连接部分上。然后,当将224*224*3图像输入模型时,它会尝试“穿过”一个线性层,其中25K个特征作为输入。要解决这个问题,您只需在之前添加卷积部分,为此,请按如下方式重新定义模型:现在,您只需告诉要优化的参数,如果您只想训练最后一个层(新添加的层),请执行以下操作:
^{pr2}$其余代码保持不变。在
希望有帮助!在
相关问题 更多 >
编程相关推荐