如何根據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