Pytorchで一部のレイヤウェイトのみをインポートする方法

通常、移行学習、すなわち比較的一般的なpretext-task上で事前訓練を行い、その後、異なるdownstream taskに対して微調整を行う.微調整の場合、ネットワーク構造の最後の数層は通常変更されます.例えば、pretext-taskはimagenet上で画像分類を行い、下流タスクは意味分割を行うと仮定すると、微調整時に分類ネットワークの最後の数層の全接続層を除去し、FCNのネットワーク構造に改造する必要がある.前のレイヤのウェイト値をロードする必要があります.
モデル構造を変更した後、torchを簡単に乱暴に使用する.load_state_dict(torch.load(‘xxx.pth’))では間違いを報告するに違いない.だから具体的にどうすればいいのか、辛抱強く下を見てください.
まず、簡単な画像分類モデルを定義します.

class model1(nn.Module):
    def __init__(self, img_size):
        super(model, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, 1, 1)
        self.conv2 = nn.Conv2d(16, 64, 3, 1, 1)
        self.fc1 = nn.Linear(self.num_feature_pixel(img_size), 1024)
        self.fc2 = nn.Linear(1024, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2,2))
        x = F.max_pool2d(F.relu(self.conv2(x)), (2,2))

        x = torch.flatten(x)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))

    def num_feature_pixel(self, img_size):
        res = 1
        for i in img_size[2:]:
            res *= i
        res = int(res * 64 / (4**2))
        return res

モデルをテストし、モデルパラメータを「pretext.pth」として保存します.

img = torch.rand([1, 3, 224, 224])
img_size = img.shape
net = model1(img_size)
res = net(img)
torch.save(net.state_dict(), 'pretext.pth')

このとき、最後の全接続レイヤを取り外してconv 3を追加すると、ネットワークの構造定義は次のようになります.

class model2(nn.Module):
    def __init__(self):
        super(model, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, 1, 1)
        self.conv2 = nn.Conv2d(16, 64, 3, 1, 1)
        self.conv3 = nn.Conv2d(64, 64, 3, 1, 1)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2,2))
        x = F.max_pool2d(F.relu(self.conv2(x)), (2,2))
        x = F.max_pool2d(F.relu(self.conv3(x)), (2,2))
        return x

このとき、新しいモデルのオブジェクトからloadを除去する前の「pretext.pth」のパラメータを実行すると、エラーが発生します.

net = model2()
net = net.load_state_dict(torch.load('pretext.pth'))

"""
RuntimeError: Error(s) in loading state_dict for model:
	Missing key(s) in state_dict: "conv3.weight", "conv3.bias". 
	Unexpected key(s) in state_dict: "fc1.weight", "fc1.bias", "fc2.weight", "fc2.bias".
"""

従来のモデルのパラメータ「pretext.pth」には、新しいモデルのconv 3パラメータは存在しないことは明らかである.同時に,fc 1とfc 2の相関パラメータも,新しいモデルにとってunexpectedである.そこで問題は,元のモデルパラメータのキーが,修正後のモデルのkeyと完全に一致しないことである.したがって、この問題を解決するには、「pretext.pth」に新しいモデルに存在するキー値ペアを抽出します.
次のコードは、問題を完全に解決します.

net = model()
pretext_model = torch.load('pretext.pth')
model2_dict = net.state_dict()
state_dict = {k:v for k,v in pretext_model.items() if k in model2_dict.keys()}
model2_dict.update(state_dict)
net.load_state_dict(model2_dict)

まずはpretext_モデルは、以前のモデルのパラメータを辞書形式で読み出し、モデル2_dictは新しいモデルを表すパラメータ辞書、state_dictは、2つのモデルが共有するパラメータキー値のペアを表します.state_を手に入れたdict以降、model 2_dictは共有するkeyを更新し,すなわち元のモデルで読めるパラメータをすべて読み込み,最後にnetをこの更新後のパラメータ辞書にロードする.