スポンサーリンク

PyTorchでモデルのパラメータ（重み）を固定する方法、転移学習やファインチューニングに利用可能

Machine Learning
2022/09/25

転移学習やファインチューニングを行う際に、モデルの重みを固定する必要があり、PyTorchにおける方法を調べてました。

モデルのパラメータ（重み）を固定する方法

PyTorchの公式Quickstartの例から、簡単なモデルを生成します。

import torch
from torch import nn

# Define model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(2, 3),
            nn.ReLU(),
            nn.Linear(3, 2),
            nn.ReLU(),
            nn.Linear(2, 1)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork()
print(model)
# NeuralNetwork(
#   (flatten): Flatten(start_dim=1, end_dim=-1)
#   (linear_relu_stack): Sequential(
#     (0): Linear(in_features=2, out_features=3, bias=True)
#     (1): ReLU()
#     (2): Linear(in_features=3, out_features=2, bias=True)
#     (3): ReLU()
#     (4): Linear(in_features=2, out_features=1, bias=True)
#   )
# )

parameters()メソッドで、重みのTensorを確認してみると、requires_grad=Trueとなっていることが分かります。requires_grad=Trueの場合は、学習時に勾配が計算され重みが更新されます。

for param in model.parameters():
  print(param)

# Parameter containing:
# tensor([[ 0.1012, -0.1174],
#         [ 0.4582,  0.5679],
#         [ 0.0129,  0.1627]], requires_grad=True)
# Parameter containing:
# tensor([ 0.3331, -0.0225,  0.4461], requires_grad=True)
# Parameter containing:
# tensor([[-0.3295, -0.1943, -0.3377],
#         [ 0.2265, -0.1137, -0.1145]], requires_grad=True)
# Parameter containing:
# tensor([0.3788, 0.3076], requires_grad=True)
# Parameter containing:
# tensor([[-0.6390, -0.3567]], requires_grad=True)
# Parameter containing:
# tensor([0.4530], requires_grad=True)

学習時に勾配が計算されないように、重みのTensorをrequires_grad=Falseにします。requires_grad=Trueと表示されていないことから、requires_gradがFalseになっていることが分かります。

for param in model.parameters():
    param.requires_grad = False
    print(param)

# Parameter containing:
# tensor([[ 0.1012, -0.1174],
#         [ 0.4582,  0.5679],
#         [ 0.0129,  0.1627]])
# Parameter containing:
# tensor([ 0.3331, -0.0225,  0.4461])
# Parameter containing:
# tensor([[-0.3295, -0.1943, -0.3377],
#         [ 0.2265, -0.1137, -0.1145]])
# Parameter containing:
# tensor([0.3788, 0.3076])
# Parameter containing:
# tensor([[-0.6390, -0.3567]])
# Parameter containing:
# tensor([0.4530])

転移学習などで最後の層のみ学習させたい場合は、パラメータを選択しrequires_grad=Trueにします。最終層の重みのみrequires_grad=Trueになっていることが確認できました。

for param in model.linear_relu_stack[4].parameters():
    param.requires_grad = True

for param in model.parameters():
    print(param)

# Parameter containing:
# tensor([[ 0.1012, -0.1174],
#         [ 0.4582,  0.5679],
#         [ 0.0129,  0.1627]])
# Parameter containing:
# tensor([ 0.3331, -0.0225,  0.4461])
# Parameter containing:
# tensor([[-0.3295, -0.1943, -0.3377],
#         [ 0.2265, -0.1137, -0.1145]])
# Parameter containing:
# tensor([0.3788, 0.3076])
# Parameter containing:
# tensor([[-0.6390, -0.3567]], requires_grad=True)
# Parameter containing:
# tensor([0.4530], requires_grad=True)

関連記事、参考資料

第11章の事前学習済みモデルの利用で、学習済みの読み込み方法やパラメータの固定方法など、転移学習/ファインチューニングに必要な内容が丁寧に説明されており、分かりやすいのでオススメです。

コメント