Google Colab で trl によるTransformerモデルの強化学習を試す
「Google Colab」で「trl」によるTransformerモデルの強化学習を試したので、まとめました。
1. trl
「trl」(Transformer Reanforcement Learning)は、強化学習でTransformerモデルを強化学習するためのパッケージです。
PPOによるTransformerモデルの強化学習は、次の3つのステップで構成されます。
2. Colabでの実行
公式サンプル「gpt2-sentiment.ipynb」を試します。
IMDBデータセットを使って、肯定的な映画レビューを生成するように「GPT-2」をファインチューニングします。肯定的な文章に報酬を与えるため、BERT分類器による感情分析を、PPOの報酬シグナルとして使用します。
(wandbとHuggingFace Hubへのpushは無効化してます)
Google Colabでの実行手順は、次のとおりです。
(1) 新規のColabのノートブックを開き、メニュー「編集 → ノートブックの設定」で「GPU」の「プレミアム」を選択。
(2) 自動リロードの設定。
# 自動リロード
%load_ext autoreload
%autoreload 2
(3) パッケージのインストール。
# パッケージのインストール
%pip install transformers trl
(4) パッケージのインポート。
# パッケージのインポート
import torch
from tqdm import tqdm
import pandas as pd
tqdm.pandas()
from transformers import pipeline, AutoTokenizer
from datasets import load_dataset
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
from trl.core import LengthSampler
(5) コンフィグの準備
# コンフィグの準備
config = PPOConfig(
model_name="lvwerra/gpt2-imdb",
learning_rate=1.41e-5,
)
sent_kwargs = {
"return_all_scores": True,
"function_to_apply": "none",
"batch_size": 16
}
「gpt2_imdb」という「GPT2」を読み込もうとしてることがわかります。このモデルは、1エポックの IMDB データセットでファインチューニングされています。。他のパラメータは、主に元の論文"Fine-Tuning Language Models from Human Preferences"から取得しています。
(6) IMDBデータセットの読み込み。
# IMDBデータセットの読み込み
def build_dataset(config, dataset_name="imdb", input_min_text_length=2, input_max_text_length=8):
tokenizer = AutoTokenizer.from_pretrained(config.model_name)
tokenizer.pad_token = tokenizer.eos_token
ds = load_dataset(dataset_name, split='train')
ds = ds.rename_columns({'text': 'review'})
ds = ds.filter(lambda x: len(x["review"])>200, batched=False)
input_size = LengthSampler(input_min_text_length, input_max_text_length)
def tokenize(sample):
sample["input_ids"] = tokenizer.encode(sample["review"])[:input_size()]
sample["query"] = tokenizer.decode(sample["input_ids"])
return sample
ds = ds.map(tokenize, batched=False)
ds.set_format(type='torch')
return ds
dataset = build_dataset(config)
def collator(data):
return dict((key, [d[key] for d in data]) for key in data[0])
IMDBデータセットには、感情を示す POSITIVE / NEGATIVE でラベルが付けられた 50,000 の映画レビューが含まれています。IMDB データセットをDataFrameに読み込み、200文字以上のコメントをフィルタで処理します。次に、各テキストをトークン化し、LengthSamplerでランダムなサイズにカットしています。
(7) モデルの読み込み。
# モデルの読み込み
model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)
tokenizer = AutoTokenizer.from_pretrained(config.model_name)
tokenizer.pad_token = tokenizer.eos_token
GPT2モデルのバリューヘッドとトークナイザーを読み込みます。 モデルは2回読み込みます。最初のモデルは最適化、2 番目のモデルは開始点からの KL 発散を計算するための参照として機能します。これは、最適化モデルが元の言語モデルから大きく逸脱しないようにするためのPPOの追加の報酬シグナルとして機能します。
(8) PPOトレーナーの準備。
# PPOトレーナーの準備
ppo_trainer = PPOTrainer(
config,
model,
ref_model,
tokenizer,
dataset=dataset,
data_collator=collator)
(9) BERT分類器の準備。
device = ppo_trainer.accelerator.device
if ppo_trainer.accelerator.num_processes == 1:
device = 0 if torch.cuda.is_available() else "cpu" # to avoid a `pipeline` bug
sentiment_pipe = pipeline("sentiment-analysis", model="lvwerra/distilbert-imdb", device=device)
モデル出力は、POSITIVE / NEGATIVEの2値分類で、POSITIVEを言語モデルの報酬シグナルとして使用します。
・POSITIVEの動作確認
text = 'this movie was really bad!!'
sentiment_pipe(text, **sent_kwargs)
[[{'label': 'NEGATIVE', 'score': 2.3350484371185303},
{'label': 'POSITIVE', 'score': -2.726576328277588}]]
・NEGATIVEの動作確認
text = 'this movie was really good!!'
sentiment_pipe(text, **sent_kwargs)
[[{'label': 'NEGATIVE', 'score': -2.294790267944336},
{'label': 'POSITIVE', 'score': 2.557040214538574}]]
(10) 生成設定の準備。
応答の生成には、サンプリングを使用するだけで、top_kとnucleus samplingとmin_lengthがオフになっていることを確認します。
gen_kwargs = {
"min_length":-1,
"top_k": 0.0,
"top_p": 1.0,
"do_sample": True,
"pad_token_id": tokenizer.eos_token_id
}
(11) モデルの最適化
A100で2時間30分ほどかかりました。
output_min_length = 4
output_max_length = 16
output_length_sampler = LengthSampler(output_min_length, output_max_length)
generation_kwargs = {
"min_length":-1,
"top_k": 0.0,
"top_p": 1.0,
"do_sample": True,
"pad_token_id": tokenizer.eos_token_id
}
for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
query_tensors = batch['input_ids']
#### GPT-2からの回答を得る
response_tensors = []
for query in query_tensors:
gen_len = output_length_sampler()
generation_kwargs["max_new_tokens"] = gen_len
response = ppo_trainer.generate(query, **generation_kwargs)
response_tensors.append(response.squeeze()[-gen_len:])
batch['response'] = [tokenizer.decode(r.squeeze()) for r in response_tensors]
#### 感情分析のスコアを計算
texts = [q + r for q,r in zip(batch['query'], batch['response'])]
pipe_outputs = sentiment_pipe(texts, **sent_kwargs)
rewards = [torch.tensor(output[1]["score"]) for output in pipe_outputs]
#### PPOのステップを実行
stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
ppo_trainer.log_stats(stats, batch, rewards)
モデルの最適化の学習ループは、次のステップで構成されます。
(12) モデルの動作確認
IMDBデータセットの例をいくつかつかって確認します。最適化前より最適化後の方がスコアが上がっていることがわかります。
#### データセットからバッチを取得
bs = 16
game_data = dict()
dataset.set_format("pandas")
df_batch = dataset[:].sample(bs)
game_data['query'] = df_batch['query'].tolist()
query_tensors = df_batch['input_ids'].tolist()
response_tensors_ref, response_tensors = [], []
#### gpt2およびgpt2_refから応答を取得
for i in range(bs):
gen_len = output_length_sampler()
output = ref_model.generate(torch.tensor(query_tensors[i]).unsqueeze(dim=0).to(device),
max_new_tokens=gen_len, **gen_kwargs).squeeze()[-gen_len:]
response_tensors_ref.append(output)
output = model.generate(torch.tensor(query_tensors[i]).unsqueeze(dim=0).to(device),
max_new_tokens=gen_len, **gen_kwargs).squeeze()[-gen_len:]
response_tensors.append(output)
#### レスポンスのデコード
game_data['response (before)'] = [tokenizer.decode(response_tensors_ref[i]) for i in range(bs)]
game_data['response (after)'] = [tokenizer.decode(response_tensors[i]) for i in range(bs)]
#### クエリ/レスポンスペアの前後における感情分析
texts = [q + r for q,r in zip(game_data['query'], game_data['response (before)'])]
game_data['rewards (before)'] = [output[1]["score"] for output in sentiment_pipe(texts, **sent_kwargs)]
texts = [q + r for q,r in zip(game_data['query'], game_data['response (after)'])]
game_data['rewards (after)'] = [output[1]["score"] for output in sentiment_pipe(texts, **sent_kwargs)]
# 結果をデータフレームに格納
df_results = pd.DataFrame(game_data)
df_results
報酬の平均/中央値を見ると、大きな違いが見られます。
print('mean:')
display(df_results[["rewards (before)", "rewards (after)"]].mean())
print()
print('median:')
display(df_results[["rewards (before)", "rewards (after)"]].median())
mean:
rewards (before) 0.090049
rewards (after) 1.672922
dtype: float64
median:
rewards (before) 0.193053
rewards (after) 1.915701
dtype: float64
(13) モデルの保存。
HuggingFace HubにはPushしない設定にしてます。
# モデルの保存
model.save_pretrained('gpt2-imdb-pos-v2', push_to_hub=False)
tokenizer.save_pretrained('gpt2-imdb-pos-v2', push_to_hub=False)
('gpt2-imdb-pos-v2/tokenizer_config.json',
'gpt2-imdb-pos-v2/special_tokens_map.json',
'gpt2-imdb-pos-v2/vocab.json',
'gpt2-imdb-pos-v2/merges.txt',
'gpt2-imdb-pos-v2/added_tokens.json',
'gpt2-imdb-pos-v2/tokenizer.json')
関連
この記事が気に入ったらサポートをしてみませんか?