2023-02-04

効率的なTransformerトレーニングのためのHugging Face Trainer

Machine Learning

NLP

Hugging Face

Python

Hugging Face Trainerクラスとは

Hugging FaceのTrainerクラスは、トランスフォーマーモデルのトレーニングとファインチューニングのプロセスを簡素化し、モデルのパフォーマンスを最適化するための直感的かつ効率的なインターフェースを提供するように設計されています。 Trainerクラスは、データ処理、最適化、および評価を含む必要なステップをカプセル化し、ユーザーが異なるモデルやハイパーパラメータの実験に集中できるようにします。

Trainerクラスには、トレーニングループのセットアップ、オプティマイザーと学習率スケジューラーの管理、トレーニングの進行状況の追跡など、一般的なタスクを処理するための様々なビルトイン関数が備わっています。さらに、カスタムトレーニングロジック、評価メトリック、およびコールバックを定義して特定のユースケースに対応することができるよう、Trainerクラスは簡単にカスタマイズ可能です。

Trainerクラスの構成要素

Trainerクラスの主要な構成要素は以下のとおりです。

モデル
トレーニングするニューラルネットワークモデルです。通常、BERT、GPT-2、RoBERTaなどのトランスフォーマーベースのアーキテクチャのインスタンスです。 Hugging Faceによって提供される事前学習済みモデルを使用するか、特定の要件に基づいて新しいモデルをインスタンス化することができます。
トレーニング引数
学習率、バッチサイズ、エポック数などのトレーニングプロセスを制御する一連のハイパーパラメータと設定です。これらの引数はTrainingArgumentsクラスのインスタンスとして、Trainerクラスに渡されます。
データセット
トレーニング、検証、テストデータを含むDatasetクラスのインスタンスです。これらのデータセットは、選択したトランスフォーマーモデルに適したトークナイザを使用して前処理とトークン化を行う必要があります。
オプティマイザーと学習率スケジューラー
デフォルトでは、TrainerクラスはAdamWオプティマイザーとウォームアップを行う線形学習率スケジューラーを使用します。ただし、カスタマイズ可能であり、独自のオプティマイザーとスケジューラーを定義し、初期化時にTrainerクラスに渡すことができます。
評価戦略
Trainerクラスは、特定のステップの後または各エポックの終了時にモデルを評価するなど、異なる評価戦略をサポートしています。 evaluation_strategy引数を使用して、好みの戦略を定義することができます。
カスタマイズオプション
Trainerクラスは高度にカスタマイズ可能で、特定のニーズに対応するために、基本的なTrainerクラスを拡張するか、特定のメソッドをオーバーライドしてカスタムトレーニングステップ、評価メトリック、およびコールバックを定義できます。

モデルのトレーニング

Trainerクラスを使用するには、Hugging Face Transformersライブラリから必要なコンポーネントをインポートし、モデルとトークナイザーをインスタンス化し、データセットを作成し、トレーニング引数を構成する必要があります。セットアップが完了したら、Trainerクラスを初期化して、そのtrain()メソッドを呼び出してトレーニングプロセスを開始します。

例えば、テキスト分類タスクのBERTモデルを微調整する場合は、次の手順を実行します。

必要なコンポーネントをインポート

python

from transformers import BertForSequenceClassification, BertTokenizerFast, Trainer, TrainingArguments

事前学習済みBERTモデルとトークナイザーをロード

python

model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')

トークナイザーを使用してデータセットを作成

python

train_dataset = ...
valid_dataset = ...

トレーニング引数を構成

python

training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
)

Trainerクラスをインスタンス化

python

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
)

モデルをトレーニング

python

trainer.train()

Trainerクラスを利用することで、テキスト分類、感情分析、質問応答など、様々な自然言語処理タスクに対して、簡単にtransformerモデルをファインチューニングできます。

新しいモデルの動的インスタンス化

Trainerクラスのmodel_initオプションを使用すると、ハイパーパラメータ検索やトレーニングループの各ランで新しいモデルを動的にインスタンス化できます。これにより、モデルの重みが各ランで再初期化され、トレーニングプロセスにクリーンなスレートが提供されます。これは、ハイパーパラメータのチューニングや、異なる構成で複数のトレーニングイテレーションを実行する場合に特に有用です。

カスタムモデルでmodel_initオプションを使用するには、独自のモデルの新しいインスタンスを初期化して返す関数を定義する必要があります。次に、この関数をインスタンス化するときにTrainerクラスに渡すことができます。

テキスト分類タスクのカスタムBERTモデルを使用してmodel_initオプションを使用する方法の例を以下に示します。

# Step 1: Import the necessary components
from transformers import BertModel, BertTokenizerFast, Trainer, TrainingArguments
import torch.nn as nn

# Step 2: Define your custom model
class CustomBertForSequenceClassification(nn.Module):
    def __init__(self, num_labels):
        super(CustomBertForSequenceClassification, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-uncased')
        self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels)

    def forward(self, input_ids, attention_mask, labels=None):
        outputs = self.bert(input_ids, attention_mask=attention_mask)
        logits = self.classifier(outputs.last_hidden_state[:, 0, :])

        if labels is not None:
            loss_fn = nn.CrossEntropyLoss()
            loss = loss_fn(logits, labels)
            return loss, logits
        else:
            return logits

# Step 3: Define a function to initialize the custom model
def model_init():
    return CustomBertForSequenceClassification(num_labels=2)

# Step 4: Load a tokenizer
tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')

# Step 5: Create the datasets using the tokenizer
train_dataset = ...
valid_dataset = ...

# Step 6: Configure the training arguments
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    evaluation_strategy="epoch",
)

# Step 7: Instantiate the Trainer class with the model_init function
trainer = Trainer(
    model_init=model_init,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
)

# Step 8: Train the custom model
trainer.train()

カスタムモデルをmodel_initオプションを使用して使用することで、各トレーニングランが新しく初期化されたインスタンスから開始されることが保証されます。これは、複数のランで一貫した結果を得たり、ハイパーパラメータチューニングを実行する場合に特に役立ちます。

カスタムトレーニングループの作成

Trainerクラスは、ほとんどのトレーニングシナリオを処理するように設計されていますが、ユーザーは独自のトレーニングループを作成して、一意の要件や制約に対応する必要がある場合があります。このセクションでは、Hugging Faceのコンポーネントを使用して、カスタムトレーニングループを作成する方法について説明します。

カスタムトレーニングループを作成するには、ユーザーは自分自身のトレーニングステップ関数を定義し、それをTrainerクラスに渡す必要があります。この関数は、モデル、データ、およびオプティマイザを入力として受け取り、順方向および逆方向のパス後に損失を出力する必要があります。ユーザーは、AutoModelおよびAutoTokenizerなど、Hugging Faceの組み込みコンポーネントを利用して入力データを処理し、モデルの更新を処理できます。

必要なコンポーネントをインポート

python

from transformers import AutoModelForSequenceClassification, AutoTokenizer, AdamW
import torch
from torch.utils.data import DataLoader

事前学習済みモデルとトークナイザーをロード

python

model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased')
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')

トークナイザを使用してデータセットを作成

python

train_dataset = ...
valid_dataset = ...

DataLoader、オプティマイザ、およびスケジューラーを設定

python

train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True)
valid_dataloader = DataLoader(valid_dataset, batch_size=64)

optimizer = AdamW(model.parameters(), lr=5e-5)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1000, gamma=0.1)

カスタムトレーニングループを定義

python

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

for epoch in range(3):
    # Training
    model.train()
    for batch in train_dataloader:
        optimizer.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        scheduler.step()

    # Validation
    model.eval()
    total_eval_loss = 0
    for batch in valid_dataloader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        with torch.no_grad():
            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss
            total_eval_loss += loss.item()

    avg_eval_loss = total_eval_loss / len(valid_dataloader)
    print(f'Epoch: {epoch + 1}, Validation Loss: {avg_eval_loss}')

この例では、3つのエポックのトレーニング、および評価フェーズが続くカスタムトレーニングループを定義しました。トレーニング中、train_dataloaderを繰り返し処理し、モデルで順方向および逆方向のパスを実行し、オプティマイザおよびスケジューラーを使用してモデルパラメータを更新します。評価中には、検証データセット上の平均損失を計算します。

カスタムトレーニングループを作成することで、トレーニングプロセスにより制御を持ち、高度なストラテジーを実装したり、デフォルトのTrainerクラスでサポートされていない特定の制約を適応することができます。