PyTorchで部分的にレイヤーをフリーズする方法｜Kaggleで使えるファインチューニング手法

Aru

この記事では、PyTorchを使って一部（特定）のレイヤー（層）をフリーズする方法を詳しく解説します。ファインチューニングでは、モデル全体ではなく部分的にレイヤーを学習させることで、精度向上や学習時間の短縮が期待できます。このテクニックは、Kaggleコンペのような場面で、既存モデルの性能を最大限に引き出したい場合に特に有効です。少しでも精度を上げたい方におすすめです。

Contents

はじめに（レイヤーフリーズの必要性）
レイヤーのフリーズ方法
transformerでEmbeddingをフリーズ
まとめ

はじめに（レイヤーフリーズの必要性）

レイヤーのフリーズとは、特定のレイヤーのパラメータ更新を停止させ、学習が行われないようにするテクニックです。

ファインチューニングでは、学習済みモデルの特定のレイヤーをフリーズして再学習を行わない設定を利用することがあります。これは、学習済みのモデルをそのまま活かしつつ、必要な部分のみを最適化するための効率的な方法です。

例えば、BERTのようなtransformerモデルでは、Embedding層が既に単語（トークン）に関して事前学習されており、特定のドメインに合わせたファインチューニングを行う際も、その層は大きく変化しないと考えられます。このため、Embedding層をフリーズすることが多いです。

もし、このような層をフリーズしなければ、少ないデータでの学習において過学習のリスクが高まり、元々持っていた汎用的な能力が低下する可能性があります。フフリーズすることで、不要な学習を抑え、元の性能を維持することができます。

一般的に、レイヤーのフリーズは、次のような状況で特に有効です。

レイヤーのフリーズが有効な例

転移学習（Transfer Learning）：事前に学習されたモデルを新しいタスクに適用する場合、モデルの一部レイヤーをフリーズして、部分的に再学習することが一般的です
過学習の防止：モデルが過学習する可能性がある場合、特に少量のトレーニングデータの場合に、一部のレイヤーをフリーズすることで、過学習を抑えることができます
計算効率の向上：大規模なモデルやデータセットの場合、すべてのレイヤーをトレーニングするのは計算上のコストが高いため、一部のレイヤーをフリーズして、計算効率を向上させることができます。transformerなどの学習の場合は、学習時間の短縮のために入力に近い層をフリーズすることもよく行います

フリーズされたレイヤーの重みは、トレーニング中に更新されません。これにより過学習の抑制や学習時間の短縮が期待できます。また、フリーズすることでGPUのメモリも節約できます。

Kaggleでは、ファインチューニングのためにレイヤーをフリーズすることがあります。特に言語モデルで多い印象です。

ここでは、Pytorchで一部レイヤーをフリーズする方法を説明します。

レイヤーのフリーズ方法

レイヤー（層）のフリーズ方法

PyTorchでレイヤーをフリーズする場合、基本的には、フリーズするレイヤーのrequires_gradをFalseに設定するだけです。

レイヤーのフリーズ方法（PyTorch）

フリーズしたいレイヤーのrequires_gradをFalseに設定

モデルのレイヤー情報を調べる方法

どのレイヤーをフリーズするか決めるためには、モデルのレイヤー構造や名前を知っておく必要があります。

例では、timmのresnet18を使います。timmでモデルを生成するコードは以下になります。

import timm

model = timm.create_model('resnet18', pretrained=True)

各層の一覧を表示するには、以下のコードを実行します。modelは、ターゲットとするモデルです。下記の例では、名前とrequires_gradの状態を表示しています。

 for n, p in model.named_parameters():
    print(n, p.requires_grad)

resnet18の場合、上記のコードを実行すると以下のような結果が出力されます。

conv1.weight True
bn1.weight True
bn1.bias True
layer1.0.conv1.weight True
layer1.0.bn1.weight True
layer1.0.bn1.bias True
layer1.0.conv2.weight True
layer1.0.bn2.weight True
layer1.0.bn2.bias True
　　　　　  :
 (途中省略）
　　　　　  :
fc.weight True
fc.bias True

特定のレイヤーのフリーズ方法

特定のレイヤーをフリーズしたい場合には、上記の名前を見ながらrequires_gradをFalseにします。例えば、conv1.weightをフリーズしたい場合には以下のようにします。

model.conv1.weight.requires_grad = False

複数のレイヤーをまとめてフリーズしたい場合

多くの場合は、いくつかのレイヤーをまとめてフリーズしたいのではないでしょうか。名前に共通の部分がある場合は、それを利用してレイヤーを特定できます。

たとえば、layer1.0という名前が含まれるレイヤーをまとめてフリーズしたい場合は、以下のように書くこともできます。

以下のコードにより、名前にlayer1.0を含むパラメータがすべてフリーズされます。

for n, p in model.named_parameters():
    if "layer1.0." in n :
      print(n)
      p.requires_grad = False

フリーズできたかどうか確認したい場合は、下記の一覧表示するコードを実行してください。

 for n, p in model.named_parameters():
    print(n, p.requires_grad)

transformerでEmbeddingをフリーズ

私の場合は、BERTなどのtransformerを利用する場合にフリーズを利用します。フリーズするのは、Embedding層と、そこから数レイヤーなことが多いです。

hugging faceで提供されているモデルの多くがEmbedding層は.embeddings.を含んでいて、エンコーダーのレイヤーはencoder.layer.0、encoder.layer.1、….という名前になっていることが多いようですので以下の関数で上位層をフリーズすることができます。

def top_n_layer_freeze(model, N):
    for n,p in model.named_parameters():
        if f".embeddings" in n:
            p.requires_grad = False

    for i in range(0,N,1):
        for n,p in model.named_parameters():
            if f'encoder.layer.{i}.' in n:
                p.requires_grad = False

例えば、上位3層をフリーズしたい場合には、以下のように呼び出します。