Hugging Faceのモデル学習で、各レイヤ毎に別のLearning Rateで学習する方法

nokomoro3

2022.11.07

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

こんちには。

データアナリティクス事業本部機械学習チームの中村です。

Hugging Faceのライブラリの使い方紹介記事第４弾です。

今回は、各レイヤ毎に別のLearning Rateで学習するためには、どのようにすればよいかを理解していきます。

前回のSchedulerと似た部分の内容になり、コードも一部流用しているため、併せて前回記事も参照ください。

層毎に別のLearning Rateで学習するとは

代表的な例としては、事前学習済みのBERT部分と、分類ヘッドで学習ヘッドを変更する例があります。

事前学習済みのBERT部分：Lerning Rateを小さく設定し、学習時に少しだけ更新する。
分類ヘッド：Lerning Rateを大きく設定し、学習時により大きく更新する。

今回は、この例にそって実装する方法を見ていきます。

実行環境

今回はGoogle Colaboratory環境で実行しました。

ハードウェアなどの情報は以下の通りです。

GPU: Tesla P100 (GPUメモリ16GB搭載)
CUDA: 11.1
メモリ: 13GB

主なライブラリのバージョンは以下となります。

transformers: 4.24.0
datasets: 2.6.1

インストール

transformersとdatasetsをインストールします。

!pip install transformers datasets

また事前学習モデルの依存モジュールをインストールします。

!pip install fugashi
!pip install ipadic
!pip install sentencepiece

ベースとするコード

今回のベースとするコードは、前回を踏襲した以下のコードとします。

49-55行目の部分がSchedulerに関する部分でした。今回はLearning Rateをすべてのstepで一定とするconstantをSchedulerにしています。

from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import TrainingArguments
from transformers import Trainer
from transformers import AdamW, get_constant_schedule
from sklearn.metrics import accuracy_score, f1_score
import torch

# データセットのロード
dataset = load_dataset("tyqiangz/multilingual-sentiments", "japanese")

# # 実験のためデータセットを縮小したい場合はコチラを有効化
# from datasets import DatasetDict
# dataset = DatasetDict({
#     "train": dataset['train'].select(range(100)),
#     "validation": dataset['validation'].select(range(100)),
#     "test": dataset['test'].select(range(100)),
# })

# トークナイザのロード
model_ckpt = "cl-tohoku/bert-base-japanese-whole-word-masking"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

# トークナイズ処理
def tokenize(batch):
    return tokenizer(batch["text"], padding=True, truncation=True)
dataset_encoded = dataset.map(tokenize, batched=True, batch_size=None)

# 事前学習モデルのロード
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_labels = 3
model = (AutoModelForSequenceClassification
    .from_pretrained(model_ckpt, num_labels=num_labels)
    .to(device))

# メトリクスの定義
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    f1 = f1_score(labels, preds, average="weighted")
    acc = accuracy_score(labels, preds)
    return {"accuracy": acc, "f1": f1}

# 学習パラメータの設定
batch_size = 16
model_name = "sample-text-classification-bert"

# modelから学習すべきパラメータを抽出
params = filter(lambda x: x.requires_grad, model.parameters())

# 今回はoptimizerにAdamWを使用
optimizer = AdamW(params, lr=2e-5)

scheduler = get_constant_schedule(optimizer)

training_args = TrainingArguments(
    output_dir=model_name,
    num_train_epochs=10,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    logging_strategy="steps",
    disable_tqdm=False,
    logging_steps=1,
    push_to_hub=False,
    log_level="error",
)

# Trainerの定義
trainer = Trainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=dataset_encoded["train"],
    eval_dataset=dataset_encoded["validation"],
    tokenizer=tokenizer,
    optimizers=[optimizer, scheduler],
)

# トレーニング実行
trainer.train()

この内容についての解説は前回の記事を参照ください。

各レイヤでLearning Rateを変更する方法

ポイントは、ベースとするコードにもある以下の部分です。

# modelから学習すべきパラメータを抽出
params = filter(lambda x: x.requires_grad, model.parameters())

ここで、今はmodel全体のパラメータを全て抽出していますが、これをレイヤごとに分けて取得できれば、異なるLearning Rateを適用することができます。

各パラメータの確認

model.parameters()では、各パラメータの名前が分かりませんが、代わりにmodel.named_parameters()を使用することでパラメータ名を確認できます。

list(name for name, param in model.named_parameters())

出力は以下となります。（長いため一部省略しています）

['bert.embeddings.word_embeddings.weight',
 'bert.embeddings.position_embeddings.weight',
 'bert.embeddings.token_type_embeddings.weight',
 'bert.embeddings.LayerNorm.weight',
 'bert.embeddings.LayerNorm.bias',
 'bert.encoder.layer.0.attention.self.query.weight',
 'bert.encoder.layer.0.attention.self.query.bias',
 'bert.encoder.layer.0.attention.self.key.weight',
 'bert.encoder.layer.0.attention.self.key.bias',
 'bert.encoder.layer.0.attention.self.value.weight',
 'bert.encoder.layer.0.attention.self.value.bias',
 'bert.encoder.layer.0.attention.output.dense.weight',
 'bert.encoder.layer.0.attention.output.dense.bias',
 'bert.encoder.layer.0.attention.output.LayerNorm.weight',
 'bert.encoder.layer.0.attention.output.LayerNorm.bias',
 'bert.encoder.layer.0.intermediate.dense.weight',
 'bert.encoder.layer.0.intermediate.dense.bias',
 'bert.encoder.layer.0.output.dense.weight',
 'bert.encoder.layer.0.output.dense.bias',
 'bert.encoder.layer.0.output.LayerNorm.weight',
 'bert.encoder.layer.0.output.LayerNorm.bias',
 'bert.encoder.layer.1.attention.self.query.weight',
 'bert.encoder.layer.1.attention.self.query.bias',
 'bert.encoder.layer.1.attention.self.key.weight',
 'bert.encoder.layer.1.attention.self.key.bias',
 'bert.encoder.layer.1.attention.self.value.weight',
 'bert.encoder.layer.1.attention.self.value.bias',
 'bert.encoder.layer.1.attention.output.dense.weight',
 'bert.encoder.layer.1.attention.output.dense.bias',
 'bert.encoder.layer.1.attention.output.LayerNorm.weight',
 'bert.encoder.layer.1.attention.output.LayerNorm.bias',
 'bert.encoder.layer.1.intermediate.dense.weight',
 'bert.encoder.layer.1.intermediate.dense.bias',
 'bert.encoder.layer.1.output.dense.weight',
 'bert.encoder.layer.1.output.dense.bias',
 'bert.encoder.layer.1.output.LayerNorm.weight',
 'bert.encoder.layer.1.output.LayerNorm.bias',
 'bert.encoder.layer.2.attention.self.query.weight',
 'bert.encoder.layer.2.attention.self.query.bias',
 'bert.encoder.layer.2.attention.self.key.weight',
 'bert.encoder.layer.2.attention.self.key.bias',
 'bert.encoder.layer.2.attention.self.value.weight',
 'bert.encoder.layer.2.attention.self.value.bias',
 'bert.encoder.layer.2.attention.output.dense.weight',
 'bert.encoder.layer.2.attention.output.dense.bias',
 'bert.encoder.layer.2.attention.output.LayerNorm.weight',
 'bert.encoder.layer.2.attention.output.LayerNorm.bias',
 'bert.encoder.layer.2.intermediate.dense.weight',
 'bert.encoder.layer.2.intermediate.dense.bias',
 'bert.encoder.layer.2.output.dense.weight',
 'bert.encoder.layer.2.output.dense.bias',
 'bert.encoder.layer.2.output.LayerNorm.weight',
 'bert.encoder.layer.2.output.LayerNorm.bias',
 
 ...中略...

 'bert.encoder.layer.11.attention.self.query.weight',
 'bert.encoder.layer.11.attention.self.query.bias',
 'bert.encoder.layer.11.attention.self.key.weight',
 'bert.encoder.layer.11.attention.self.key.bias',
 'bert.encoder.layer.11.attention.self.value.weight',
 'bert.encoder.layer.11.attention.self.value.bias',
 'bert.encoder.layer.11.attention.output.dense.weight',
 'bert.encoder.layer.11.attention.output.dense.bias',
 'bert.encoder.layer.11.attention.output.LayerNorm.weight',
 'bert.encoder.layer.11.attention.output.LayerNorm.bias',
 'bert.encoder.layer.11.intermediate.dense.weight',
 'bert.encoder.layer.11.intermediate.dense.bias',
 'bert.encoder.layer.11.output.dense.weight',
 'bert.encoder.layer.11.output.dense.bias',
 'bert.encoder.layer.11.output.LayerNorm.weight',
 'bert.encoder.layer.11.output.LayerNorm.bias',
 'bert.pooler.dense.weight',
 'bert.pooler.dense.bias',
 'classifier.weight',
 'classifier.bias']

簡単に説明すると、以下の部分が先頭の各種埋め込み層とLayerNormになります。

['bert.embeddings.word_embeddings.weight',
 'bert.embeddings.position_embeddings.weight',
 'bert.embeddings.token_type_embeddings.weight',
 'bert.embeddings.LayerNorm.weight',
 'bert.embeddings.LayerNorm.bias',

以下はEncoderの１層分です。baseではこれが12層あるため、layer.0～layer.11まであります。

 'bert.encoder.layer.0.attention.self.query.weight',
 'bert.encoder.layer.0.attention.self.query.bias',
 'bert.encoder.layer.0.attention.self.key.weight',
 'bert.encoder.layer.0.attention.self.key.bias',
 'bert.encoder.layer.0.attention.self.value.weight',
 'bert.encoder.layer.0.attention.self.value.bias',
 'bert.encoder.layer.0.attention.output.dense.weight',
 'bert.encoder.layer.0.attention.output.dense.bias',
 'bert.encoder.layer.0.attention.output.LayerNorm.weight',
 'bert.encoder.layer.0.attention.output.LayerNorm.bias',
 'bert.encoder.layer.0.intermediate.dense.weight',
 'bert.encoder.layer.0.intermediate.dense.bias',
 'bert.encoder.layer.0.output.dense.weight',
 'bert.encoder.layer.0.output.dense.bias',
 'bert.encoder.layer.0.output.LayerNorm.weight',
 'bert.encoder.layer.0.output.LayerNorm.bias',

以下はEncoderの最終層の先頭(CLSトークン相当部分)に線形層を通すpoolerの部分です。

 'bert.pooler.dense.weight',
 'bert.pooler.dense.bias',

最後に分類のためのヘッド(線形層)があります。

 'classifier.weight',
 'classifier.bias']

構造の詳細は、BERTの元論文やTransformerの解説をご覧ください。

設定方法

事前学習モデルの部分には"bert"という文字列、ヘッドには"classifier"という文字列で見分けて設定します。

body_params \
    = [p for n,p in model.named_parameters() if n.startswith("bert") and p.requires_grad]
head_params \
    = [p for n,p in model.named_parameters() if n.startswith("classifier") and p.requires_grad]

# 別々にLearning Rateを設定するには以下のようにする
optimizer = AdamW([
    {"params": body_params, 'lr': 0.0}, 
    {"params": head_params, 'lr': 2e-5}
])

scheduler = get_constant_schedule(optimizer)

動作確認

動作確認をします。念のため、先頭と末尾のパラメータをバックアップして学習前後の変化を確認します。

from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import TrainingArguments
from transformers import Trainer
from transformers import AdamW, get_constant_schedule
from sklearn.metrics import accuracy_score, f1_score
import torch

# データセットのロード
dataset = load_dataset("tyqiangz/multilingual-sentiments", "japanese")

# # 実験のためデータセットを縮小したい場合はコチラを有効化
# from datasets import DatasetDict
# dataset = DatasetDict({
#     "train": dataset['train'].select(range(100)),
#     "validation": dataset['validation'].select(range(100)),
#     "test": dataset['test'].select(range(100)),
# })

# トークナイザのロード
model_ckpt = "cl-tohoku/bert-base-japanese-whole-word-masking"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

# トークナイズ処理
def tokenize(batch):
    return tokenizer(batch["text"], padding=True, truncation=True)
dataset_encoded = dataset.map(tokenize, batched=True, batch_size=None)

# 事前学習モデルのロード
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_labels = 3
model = (AutoModelForSequenceClassification
    .from_pretrained(model_ckpt, num_labels=num_labels)
    .to(device))

backup1 = list(model.parameters())[0].to('cpu').detach().numpy().copy()
backup2 = list(model.parameters())[-1].to('cpu').detach().numpy().copy()

# メトリクスの定義
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    f1 = f1_score(labels, preds, average="weighted")
    acc = accuracy_score(labels, preds)
    return {"accuracy": acc, "f1": f1}

# 学習パラメータの設定
batch_size = 16
model_name = "sample-text-classification-bert"

# 事前学習モデルの部分には"bert"という文字列、ヘッドには"classifier"という文字列で見分けて設定
body_params = [p for n,p in model.named_parameters() if n.startswith("bert") and p.requires_grad]
head_params = [p for n,p in model.named_parameters() if n.startswith("classifier") and p.requires_grad]

# 別々にLearning Rateを設定するには以下のようにする
optimizer = AdamW([
    {"params": body_params, 'lr': 0}, 
    {"params": head_params, 'lr': 2e-5}
])

scheduler = get_constant_schedule(optimizer)

training_args = TrainingArguments(
    output_dir=model_name,
    num_train_epochs=10,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    logging_strategy="steps",
    disable_tqdm=False,
    logging_steps=1,
    push_to_hub=False,
    log_level="error",
)

# Trainerの定義
trainer = Trainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=dataset_encoded["train"],
    eval_dataset=dataset_encoded["validation"],
    tokenizer=tokenizer,
    optimizers=[optimizer, scheduler],
)

# トレーニング実行
trainer.train()

以下のように、先頭の層は値が変わらないことを確認できました。

backup1 - list(model.parameters())[0].to('cpu').detach().numpy().copy()

array([[0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.],
       ...,
       [0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.]], dtype=float32)

また以下のように、末尾の層はパラメータが更新されていることが分かります。

backup2 - list(model.parameters())[-1].to('cpu').detach().numpy().copy()

array([ 0.00139277,  0.00139004, -0.00139202], dtype=float32)

今回のように複数のLearning Rateを設定した場合、trainerのログには先頭のLearning Rateが記載されるようですので、ご注意ください。

train_log = [i for i in trainer.state.log_history if "loss" in i]
train_log

[{'loss': 0.9356, 'learning_rate': 0, 'epoch': 0.14, 'step': 1},
 {'loss': 0.9791, 'learning_rate': 0, 'epoch': 0.29, 'step': 2},
 {'loss': 0.9827, 'learning_rate': 0, 'epoch': 0.43, 'step': 3},
 {'loss': 0.9548, 'learning_rate': 0, 'epoch': 0.57, 'step': 4},
 {'loss': 0.8873, 'learning_rate': 0, 'epoch': 0.71, 'step': 5},
 {'loss': 0.9497, 'learning_rate': 0, 'epoch': 0.86, 'step': 6},
 {'loss': 0.844, 'learning_rate': 0, 'epoch': 1.0, 'step': 7},
 {'loss': 0.903, 'learning_rate': 0, 'epoch': 1.14, 'step': 8},
 {'loss': 0.9116, 'learning_rate': 0, 'epoch': 1.29, 'step': 9},
 {'loss': 0.8593, 'learning_rate': 0, 'epoch': 1.43, 'step': 10},
 {'loss': 0.9059, 'learning_rate': 0, 'epoch': 1.57, 'step': 11},
 {'loss': 0.8507, 'learning_rate': 0, 'epoch': 1.71, 'step': 12},
 {'loss': 0.8737, 'learning_rate': 0, 'epoch': 1.86, 'step': 13},
 ...

より進んだ応用例

より進んだ応用例としては以下も参考にされてください。weight_decayなども各層で変更している例となります。

まとめ

いかがでしたでしょうか？

実際にBERTなどのfine-tuningは、より大規模モデル(largeなど)になるにつれて学習が難しくなるタスクもあるようです。そういった際に、Learning Rateの細かい調整が必要になるシーンがあるため記事にしてみました。

本記事がHugging Faceを使われる方の参考になれば幸いです。

Hugging Faceのモデル学習で、各レイヤ毎に別のLearning Rateで学習する方法

層毎に別のLearning Rateで学習するとは

実行環境

インストール

ベースとするコード

各レイヤでLearning Rateを変更する方法

各パラメータの確認

設定方法

動作確認

より進んだ応用例

まとめ

イベント

EVENT【4/23リモート】クラスメソッドのフリーランスエンジニア会社説明会〜AWS / システムアーキテクト案件特集〜を開催します

EVENT【5/28（火）】AWSを最大活用するための1dayカンファレンス

EVENT【5/17（金）】認証機能の開発工数削減をデモで体験！次世代認証基盤サービス『Auth0 by Okta』導入実践ウェビナー

EVENT【5/14（火）】アノテーションの中途向けオンライン会社説明会を開催します

EVENT【5/23（木）】ユースケースに学ぶAWS運用のノウハウ～可視化から統制まで～

EVENT【5/16（木）】Snowflakeを触ってみよう！初めての方向けハンズオンセミナー

EVENT【5/10（金）名古屋】クラスメソッドグループの会社説明会を開催します！

EVENT【4/25（木）リモート】Google Cloudに携わるエンジニアのキャリア～クラスメソッド会社説明会～

EVENT【4/30（火）リモート】クラスメソッドの会社説明会を開催します

EVENT【5/9（木）】組織的なクラウド統制はじめの一歩

Hugging Faceのモデル学習で、各レイヤ毎に別のLearning Rateで学習する方法

層毎に別のLearning Rateで学習するとは

実行環境

インストール

ベースとするコード

各レイヤでLearning Rateを変更する方法

各パラメータの確認

設定方法

動作確認

より進んだ応用例

まとめ

イベント

EVENT【4/23リモート】クラスメソッドのフリーランスエンジニア会社説明会 〜AWS / システムアーキテクト案件特集〜 を開催します

EVENT【5/28（火）】AWSを最大活用するための1dayカンファレンス

EVENT【5/17（金）】認証機能の開発工数削減をデモで体験！次世代認証基盤サービス『Auth0 by Okta』導入実践ウェビナー

EVENT【5/14（火）】アノテーションの中途向けオンライン会社説明会を開催します

EVENT【5/23（木）】ユースケースに学ぶAWS運用のノウハウ～可視化から統制まで～

EVENT【5/16（木）】Snowflakeを触ってみよう！初めての方向けハンズオンセミナー

EVENT【5/10（金） 名古屋】クラスメソッドグループの会社説明会を開催します！

EVENT【4/25（木）リモート】Google Cloudに携わるエンジニアのキャリア～クラスメソッド会社説明会～

EVENT【4/30（火）リモート】クラスメソッドの会社説明会を開催します

EVENT【5/9（木）】組織的なクラウド統制 はじめの一歩

関連記事

PythonでCron式から実行日付を取得する

How to Fetch Latest Objects from AWS S3 and Store Locally using Python

WindowsマシンでPython環境構築をpoetry + pyenvで再度やってみた

[Amazon Bedrock][LangChain] チャット会話履歴をセッション毎に記憶・保存する方法

EVENT【4/23リモート】クラスメソッドのフリーランスエンジニア会社説明会〜AWS / システムアーキテクト案件特集〜を開催します

EVENT【5/10（金）名古屋】クラスメソッドグループの会社説明会を開催します！

EVENT【5/9（木）】組織的なクラウド統制はじめの一歩