🐙

TransformersのTrainerを汎用的に使用する

Daisuke Makino

2023/12/21に公開

この記事はJij Inc. Advent Calendar 2023の21日目の記事です。
はじめまして、株式会社Jijの牧野です。

PyTorchのモデルを訓練させる時に、Trainer系のライブラリを使ってコードをできるだけシンプルにしたい、楽をしたいという人は多いと思います。多くのTrainer系のライブラリがあって私自身も色々試してきましたが、あれこれやっている内に結局自分でモデル専用の訓練コードを書くということを繰り返しています。ここ数ヶ月研究や遊びでLLMを触ることが多くなってきて、transformersのTrainerを使うことも多くなりました。一般的なPyTorchのモデルでも普通に使えて感触も良いので、今後はtransformersのTrainerを使用して冗長な訓練コードと（ようやく）おさらばしようと思っています。今回はHugging Faceが提供するtransformersのTrainerで一般的なPyTorchのモデルを学習させる基本的な方法を紹介しようと思います。

transformers Trainerのメリット・デメリット

基本的な機能（ロギング、モデルの自動保存、カスタムトレーニングループの定義など）は当然として、transformersのTrainerを使うことのメリットは以下であると考えいます。

transformersに実装されている多様なLLMの学習をすぐに実行できる
WandB, MLflow, TensorBoardなど実験管理ツールと連携が容易
DeepSpeed対応
アップデートが早い

逆にデメリットしては、

Datasetをサブラクス化して学習データを作成するとき、書き方が少し特殊
多くのオプションがあるので覚えるのが大変
柔軟にカスタムできるが、意外とドキュメントがないので結構ソースを見に行くことになる

などが挙げられると思います。（あくまで主観です）

基本的な使い方

transformersといえばLLMですが、今回はTrainerの振る舞いに焦点を当てたいので、学習対象のモデルと訓練データを適当に用意します。まずはコードの全体像を見てみましょう。
transformersのバージョンは4.34.1です。

import torch
import torch.nn as nn

from transformers import Trainer, TrainingArguments
from torch.utils.data import Dataset


# とりあえずのモデル
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Sequential(
            nn.Linear(in_features=64, out_features=32),
            nn.ReLU(),
            nn.Linear(in_features=32, out_features=16),
            nn.ReLU(),
            nn.Linear(in_features=16, out_features=2),
        )

    def forward(self, x):
        return self.linear(x)


# とりあえずのデータセット
class SampleDataset(Dataset):
    def __init__(self, size=1000):
        self.inputs = torch.randn(size, 64)
        self.labels = torch.randint(0, 2, (size,))

    def __len__(self):
        return len(self.inputs)

    def __getitem__(self, idx):
        return {"x": self.inputs[idx], "labels": self.labels[idx]}


# 学習の設定
training_args = TrainingArguments(output_dir="./results")

model = Net()
train_dataset = SampleDataset(size=1000)
eval_dataset = SampleDataset(size=100)

# 学習の実行
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()

これだけで最低限動きますが、学習の設定にはどのようなオプションがあるのか、forward計算に渡す引数とユーザー定義のDataset（今回はSampleDataset）が返す値の関係、ロスをカスタムする方法など気になる点がいくつかあるので一つずつ見ていきます。

学習の設定

主要なTrainingArgumentsのオプションを以下に示します。

オプション	概要
`output_dir`	トレーニング結果（モデルのチェックポイント、設定など）を保存するディレクトリパス
`num_train_epochs`	トレーニングのエポック数
`do_train`	トレーニングを実行するかどうか
`do_eval`	検証セットでの評価を実行するかどうか
`per_device_train_batch_size`	各デバイス（GPU/TPU/CPU）でのトレーニング時のバッチサイズ
`per_device_eval_batch_size`	各デバイス（GPU/TPU/CPU）での評価時のバッチサイズ
`learning_rate`	トレーニングの初期学習率
`warmup_steps`	学習率が最大値に到達するまでのステップ数
`weight_decay`	重み減衰（L2正則化）の強度
`evaluation_strategy`	トレーニング中の評価のタイミングを設定する。 "no": トレーニング中に評価は行われない。 "steps": 一定のバッチ処理毎に評価を行う。具体的なステップ数は`eval_steps`で指定する。 "epoch": 各エポック毎に評価が行われる。
`eval_steps`	`evaluation_strategy`で"steps"が選択された時のステップ数
`logging_strategy`	トレーニング中のログ出力するタイミングを設定する。 "no": トレーニング中に評価は行われない。 "steps":一定のバッチ処理毎にログ出力する。具体的なステップ数は`logging_steps`で指定する。 "epoch": 各エポック毎にログ出力する。
`logging_steps`	`logging_strategy`で"steps"が選択された時のステップ数
`save_steps`	チェックポイントが更新されるまでのステップ数

他のオプションについては、こちらを参照してください。

Datasetの書き方

SampleDatasetを見ると__getitem__のリターンが辞書になっています。これはtransformersのTrainerを使う時のお作法の一つになっています。この辞書には、モデルのforwardメソッドの引数名をキーに持つ値をセットしなければなりません。今回の例で言うとforwardの引数にxがありますので、__getitem__が返す辞書は{"x": ..., }となっている必要があります。forwardに引数を追加した場合には、当然ですがDatasetにも対応するキーバリューを入れる必要があります。
では、labelsはどこで使われているのでしょうか？これは、Trainerに標準実装されているcompute_lossメソッドで使われています。デフォルトのcompute_lossはDatasetが返す辞書のlabelsキーに入っている値を正解データとして扱います。ロス関数はラベルスムージングを適用したNLLロスとして実装されています。
TrainingArgumentsを以下の様に設定するとロスの経過がエポック毎に見えるようになると思います。

training_args = TrainingArguments(
    output_dir='./results',
    logging_strategy="epoch"
)

カスタムロスの書き方

ユーザー定義のロス関数を書きたい場合には、Trainerクラスをサブクラス化し、compute_lossメソッドをオーバライドする事で実現できます。

class MyTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")
        logits = model(**inputs)
        loss = F.cross_entropy(logits, labels)
        return (loss, outputs) if return_outputs else loss

基本的にはこれだけでOKです。model引数にはTrainerに渡したモデルが、inputsにはDatasetが返す辞書が入っていますので、inputsを通してロス関数内で入力データを制御することができます。

まとめ

今回はtransformersのTrainerを使って学習をカスタムする基本的な方法を紹介しました。transformersには他にも、DeepSpeedとの連携、MLflow, WnadB, TensorBoardなどの実験管理ツールとの連携など様々なオプションがあります。これら発展的な機能については次回以降紹介していきたいと思います。

最後に

＼Rustエンジニア・数理最適化エンジニア募集中！／
株式会社Jijでは、数学や物理学のバックグラウンドを活かし、量子計算と数理最適化のフロンティアで活躍するRustエンジニア、数理最適化エンジニアを募集しています！
詳細は下記のリンクからご覧ください。皆さんのご応募をお待ちしております！
Rustエンジニア: https://open.talentio.com/r/1/c/j-ij.com/pages/51062
数理最適化エンジニア: https://open.talentio.com/r/1/c/j-ij.com/pages/75132

Jij Tech BlogPublication

transformers Trainerのメリット・デメリット

基本的な使い方

学習の設定

Datasetの書き方

カスタムロスの書き方

まとめ

最後に

Discussion