torchvision.modelsを転移学習するとき出力次元を変更する方法

カスタマイズ

pytorchでtorchvision.modelsの学習済みモデルで転移学習を試すとき、
最終層の出力次元をモデルに応じて変更させる方法。

import torch
import torchvision


def customize(model, out_features):
    """
    Parameters
    ----
    model:
        Model instance. (from torchvision.models.*)
    out_featurs:
        New dimension of out_features.
    """
    name, last_child = list(model.named_children())[-1]
    if type(last_child) == torch.nn.modules.Sequential:
        args = list(last_child)
        device = next(args[-1].parameters()).device
        in_features = args[-1].in_features
        args[-1] = torch.nn.Linear(in_features, out_features)
        modified = torch.nn.modules.Sequential(*args)
    if type(last_child) == torch.nn.modules.linear.Linear:
        device = next(last_child.parameters()).device
        in_features = last_child.in_features
        modified = torch.nn.Linear(in_features, out_features)
    setattr(model, name, modified.to(device))


model = torchvision.models.vgg16(pretrained=True)
customize(model, 2)
print(model(torch.zeros(1, 3, 224, 224)).shape) # torch.Size([1, 2])


model = torchvision.models.densenet169(pretrained=True)
customize(model, 2)
print(model(torch.zeros(1, 3, 224, 224)).shape) # torch.Size([1, 2])

#以下覚書

コーディングでハマった所メモ。
model.children()[-1]等でモデルの最後の子にアクセスできるが、
モデルごとにSequentialの場合とLinearの場合が存在することと、(例: vgg16とdensenet169)
model.classifier = Linear(in_features, out_features)
等の様にmodelのアトリビュートからアクセスするなら上手くいくが、
model.children()[-1] = Linear(in_features, out_features)
等の様に取得した子を上書きしても上手く更新できない。
モデルごとに最終層の名称がclassifierだったりfcだったりと異なるので、
named_children()から名前を取得した上でsetattr()を使うと上手く更新できた。