如何根据torch.model修改最终的FC层

 

我想在我自己的数据集中使用VGG19，它有8个类。因此，我想将最后一个fc层的输出更改为8.所以我应该怎么做来改变最后一个fc层来适应它。

 

model = torchvision.models.vgg19(pretrained=True)
for param in model.parameters():
    param.requires_grad = False
    # Replace the last fully-connected layer
    # Parameters of newly constructed modules have requires_grad=True by default
model.fc = nn.Linear(512, 8) # assuming that the fc7 layer has 512 neurons, otherwise change it 
model.cuda()

 

 

Reference

1  https://discuss.ptorch.com/question/95