ファインチューニングで一部のレイヤー（層）をフリーズする方法 | Kaggle TIPS

ファインチューニングを行う場合、一部のレイヤーをフリーズすることで性能が向上することがあります。この記事では、PyTorchでレイヤーをフリーズする方法について解説します。少しでも性能を上げたい場合に有用なテクニックです。

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

ファインチューニングでは、特定のレイヤー（層）をフリーズして学習させないようにする場合があります（学習済みのモデルをファインチューニングする場合などに使用されるテクニックです）。

例えば、BERTなどのtransformerの場合、Embedding層などは単語（トークン）に関しては事前学習しているため、ドメインに合わせてファインチューニングする場合であってもあまり変化しないと考えられます。

このような層を、少ないデータでファインチューニングすると逆に過学習してしまう可能性もあるため、学習しないように設定した方が結果として性能が向上することもあります。

レイヤーのフリーズは、次のような場合に有用です。

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

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

ここでは、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層をフリーズしたい場合には、以下のように呼び出します。

top_n_layer_freeze(model, 3)

まとめ

transformerのファインチューニングではEmbedding層＋数層をフリーズすることで精度を向上できることがあります。また、フリーズによりGPUのメモリの節約もできます。特に、大きなモデルの場合、再学習の必要のないレイヤーをフリーズするのは学習時間の短縮に有用です。