PyTorchでFine-Tuningをするためにモデルをカスタムするときの注意

VGGシリーズのように構築済みのモデルをFine-Tuningするとき、出力サイズに合わせてFull-Connectionのパラメータだけ変えようとしてハマったのでメモ

VGG-NETS

pytorch.org

問題

例えばVGG19で、最終FCの出力を1000から10にするために

model = torch.hub.load('pytorch/vision:v0.6.0', 'vgg19', pretrained=True)
model.classifier[6].out_features = 10

このように指定するとmodelのprint結果自体は正しく10になっている

print(model)
>
...

(2): Dropout(p=0.5, inplace=False)
(3): Linear(in_features=4096, out_features=4096, bias=True)
(4): ReLU(inplace=True)
(5): Dropout(p=0.5, inplace=False)
(6): Linear(in_features=4096, out_features=10, bias=True)
...

しかしこれで実際にモデルに画像を入れてみると、出力が1000のまま変わってない

outputs = net(images)
print(outputs.size())

>torch.Size([64, 1000]) #バッチサイズ64

解決策

FCの出力をちゃんと10にするにはこうしなければいけない

model.classifier[6] = nn.Linear(4096, 10)

outputs = net(images)
print(outputs.size())

>torch.Size([64, 10]) #バッチサイズ64

参考

How to modify the final FC layer based on the torch.model - PyTorch Forums

Transfer Learning with PyTorch : Learn to Use Pretrained VGG16 Model 間違っている方の例