Huggingface Transformers 入門 (4) - 訓練とファインチューニング

以下の記事を参考に書いてます。

・Huggingface Transformers : Training and fine-tuning

前回

1. PyTorchでのファインチューニング

「TF」で始まらない「Huggingface Transformers」のモデルクラスはPyTorchモジュールです。推論と最適化の両方でPyTorchのモデルと同じように利用できます。

テキスト分類のデータセットでモデルをファインチューニングする一般的なタスクを考えてみます。from_pretrained()を用いてモデルをインスタンス化すると、指定されたモデルの「モデルの構成」と「事前学習した重み」が、モデルの初期化に使用されます。このライブラリには，指定された事前学習済みモデルに含まれていない場合には、ランダムにインスタンス化される重みを持つタスク固有の「最終層」または「ヘッド」も多数含まれています。例えば、BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2) でモデルをインスタンス化すると、「bert-base-uncased」モデルからコピーされたエンコーダ重みと、出力サイズ2のエンコーダの上にあるランダムに初期化されたテキスト分類ヘッドを持つ モデルのインスタンスが作成されます。

モデルは、デフォルトでは「評価モード」で初期化されます。model.train() を呼び出すことで、「訓練モード」にすることができます。

from transformers import BertForSequenceClassification
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', return_dict=True)
model.train() # 訓練モードに変更

これは、事前学習したBERTエンコーダを利用して、任意のテキスト分類のデータセットで簡単に学習できるので便利です。PyTorchオプティマイザを使用することができますが、私たちのライブラリには、「重み減衰」と同様に「勾配バイアス補正」を実装したAdamWオプティマイザも用意されています。

from transformers import AdamW
optimizer = AdamW(model.parameters(), lr=1e-5) # AdamWオプティマイザ

オプティマイザでは、特定のパラメータグループに異なるハイパーパラメータを適用することができます。例えば、バイアスとレイヤ正規化の項以外のすべてのパラメータに「重み減衰」を適用することができます。

no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
    {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
    {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
optimizer = AdamW(optimizer_grouped_parameters, lr=1e-5) # オプティマイザ 

ここで、__call__()を使って簡単なダミーの学習バッチをセットアップします。これはBatchEncoding()インスタンスを返し、モデルに渡す必要があるかもしれないもの全てを準備します。

from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
text_batch = ["I love Pixar.", "I don't care for Pixar."]
encoding = tokenizer(text_batch, return_tensors='pt', padding=True, truncation=True)
input_ids = encoding['input_ids']
attention_mask = encoding['attention_mask']

labels引数を指定して分類モデルを呼び出すと、最初に返される要素は、予測値と渡されたラベルの間のクロスエントロピー損失です。オプティマイザの設定が済んでいるので、逆パスを行い、重みを更新することができます。

labels = torch.tensor([1,0]).unsqueeze(0)
outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs.loss
loss.backward()
optimizer.step()

あるいは、ロジットを取得して自分で損切りを計算すればいいだけです。
以下は先ほどの例と同等です。

from torch.nn import functional as F
labels = torch.tensor([1,0]).unsqueeze(0)
outputs = model(input_ids, attention_mask=attention_mask)
loss = F.cross_entropy(labels, outputs.logitd)
loss.backward()
optimizer.step()

もちろん、通常通りモデルと入力に対して、to('cuda')を呼び出すことでGPU上での学習も可能です。

また、学習率スケジューリングツールもいくつか用意しています。以下のようにすれば、num_warmup_steps分だけウォーミングアップして、訓練終了までに0に線形に減衰するスケジューラを設定することができます。

from transformers import get_linear_schedule_with_warmup
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_train_steps)

あとは、オプティマイザ.step()の後にscheduler.step()を呼び出すだけです。

loss.backward()
optimizer.step()
scheduler.step()

後述するTrainer()を使うことを強くお勧めします。これは、混合精度や簡単なテンソルボードロギングのような特徴を持つ「Huggingface Transformers」モデルの訓練の可動部分を便利に処理します。

◎ エンコーダの凍結
場合によっては、事前学習したエンコーダの重みを凍結したままにして、ヘッド層の重みだけを最適化したいと思うかもしれません。そのためには、エンコーダパラメータのrequires_grad属性をFalseに設定するだけです。

for param in model.base_model.parameters():
    param.requires_grad = False

2. TensorFlow 2でのファインチューニング

モデルは「TensorFlow 2」で学習することもできます。「PyTorch」と同様に、TensorFlowモデルはfrom_pretrained()でインスタンス化して、事前学習されたモデルからエンコーダの重みをロードすることができます。

from transformers import TFBertForSequenceClassification
model = TFBertForSequenceClassification.from_pretrained('bert-base-uncased')

tensorflow_datasetsを使って、GLUEからMRPCデータセットをロードしてみましょう。次に、組み込みのglue_convert_examples_to_features()を使って MRPCをトークン化し、TensorFlow Datasetオブジェクトに変換します。トークナイザーはフレームワークに依存しないので、事前学習されたトークナイザーの前にTFを付ける必要はないことに注意してください。

from transformers import BertTokenizer, glue_convert_examples_to_features
import tensorflow as tf
import tensorflow_datasets as tfds
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
data = tfds.load('glue/mrpc')
train_dataset = glue_convert_examples_to_features(data['train'], tokenizer, max_length=128, task='mrpc')
train_dataset = train_dataset.shuffle(100).batch(32).repeat(2)

その後、モデルをコンパイルし、他のKerasモデルとして学習することができます。

optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer=optimizer, loss=loss)
model.fit(train_dataset, epochs=2, steps_per_epoch=115)

「TensorFlowとPyTorchモデル間の相互運用性」により、モデルを保存してからPyTorchモデルとしてリロード（またはその逆）することもできます。

from transformers import BertForSequenceClassification
model.save_pretrained('./my_mrpc_model/')
pytorch_model = BertForSequenceClassification.from_pretrained('./my_mrpc_model/', from_tf=True)

3. Trainer

Trainer(TFTrainer)は、シンプルながらも機能が充実した訓練と評価のインターフェースを提供しています。訓練オプションと、ロギング、勾配蓄積、混合精度などの組み込み機能を使って、モデルの訓練、ファインチューニング、評価を行うことができます。

from transformers import BertForSequenceClassification, Trainer, TrainingArguments

# モデルの準備
model = BertForSequenceClassification.from_pretrained("bert-large-uncased")

# Trainerのパラメータの準備
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の準備
trainer = Trainer(
    model=model,                     # モデル
    args=training_args,              # 訓練引数
    train_dataset=train_dataset,     # 訓練データセット
    eval_dataset=test_dataset        # 評価データセット
)

ここで、trainer.train()を呼び出して訓練を行い、trainer.evaluate()を呼び出して評価を行います。独自のモジュールを使用することもできますが、forwardから返される最初の引数は最適化したい損失でなければなりません。

Trainer()は、バッチを照合してモデルに投入する準備をするために、組み込みのデフォルト関数を使用します。必要であれば、data_collator引数を使用して、データセットが提供するフォーマットでデータを取り込み、モデルに投入する準備ができたバッチを返す独自の照合関数を渡すこともできます。TFTrainer()は、渡されたデータセットがtensorflow_datasetsのデータセットオブジェクトであることを期待していることに注意してください。

損失に加えて追加のメトリクスを計算するために、独自のcompute_metrics関数を定義してトレーナーに渡すこともできます。

from sklearn.metrics import accuracy_score, precision_recall_fscore_support

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='binary')
    acc = accuracy_score(labels, preds)
    return {
        'accuracy': acc,
        'f1': f1,
        'precision': precision,
        'recall': recall
    }

最後に、指定したlogging_dirフォルダでtensorboardを起動して、計算されたメトリクスを含む結果を見ることができます。

4. 追加リソース

・Trainerを使った感情分類のColabデモ
・Huggingface Transformersのサンプル
・Trainerを使った言語モデリングのColabデモ
・様々なタスクのノートが収録されているColabデモ

次回

BERT/GPT-3/DALL-E 自然言語処理・画像処理・音声処理 人工知能プログラミング実践入門