はじめての自然言語処理

前回が分量的にやたらと重かったので、今回はその反省(反動?)を踏まえて軽い感じでいってみます。第7回で紹介した T5 ですが Hugging Face の Transformers でもサポートされてますので、その使用方法をご紹介したいと思います。

1. はじめに

今回は久しぶりに T5 の話です。T5 に関しては第7回、第8回で一度紹介しているので、未読の方は記事に目を通してから戻ってきて頂けると、より理解がしやすいと思います。

さて、 T5 ですが Google のオリジナルコード(以下 “t5"と記述)¹は敷居が高いと感じる方もいらっしゃるのではないでしょうか。 Estimator API ベースのコードや gin による設定など慣れていないと、とっつきにくいのではないかと思います。

そこで今回は Hugging Face の Transformers ²を使って T5 を動かす方法をご紹介します。

Transformers は BERT, GPT-2, XLNet 等々の Transformer ベースのモデルを簡単に利用することが出来るライブラリです。 ちなみに T5 は 2.3.0 でサポートされました³。 こちらの記事⁴によると FP16 での動作もサポートされたとのことで、記事中で Hugging Face の中の人曰く "T5 is one most widely used model in the library” とのことです。

それでは実際に PyTorch で T5 を動かしてみましょう! と言いたいところですが、日本語で動かすとなると事前学習モデルが必要ですよね。 t5 はちょこちょこバージョンが上がって第7回で紹介したコードのままでは動かなくなっているので⁵、改めて執筆時点の最新版で事前学習済み モデルを作るところから紹介したいと思います。

2021/4/12 追記

本記事を書きあげた後で気づいたのですが、T5 の日本語事前学習済みモデルを @sonoisa さんが公開して下さっています⁶! なので、手っ取り早く試したい人は次章をすっ飛ばして、 @sonoisa さんのモデルを使っていただくとよろしいんじゃないかと思います。

2. Google のコードによる事前学習

それでは t5 を使った事前学習からはじめていきましょう。 記事の趣旨からいうと事前学習から Transformers でやるべきな気もしますが、事前学習データの加工処理等が面倒なので、

1. t5 で事前学習済みモデルを作成
2. 1. のチェックポイントを Transformers でインポート

という段取りにしたいと思います⁷。

いつものように、記事内のコードスニペットは、特に断りがない場合は Google Colaboratory (以下、Colab)で動かす想定にしています。ノートブックを開き、アクセラレータは TPU を選んで下さい。

チェックポイントは GCS に書き込むので認証を通しておきましょう。

from google.colab import auth
auth.authenticate_user()

Tensorflow は 2.x を選択しました。最新版は 2.x でも動作するようになっています。

%tensorflow_version 2.x 

以下のようにしてインストールします。執筆時点では t5 が 0.9.0 、mesh-tensorflow が 0.1.18 でした。

!pip install t5[gcp]

ここから日本語での事前学習タスクを登録します。まずは SentencePiece のモデルを以下のようにして上書きます。 ここで gs://somewhere/t5/sentencepiece/wikipedia_20190301_ja_v003.model は日本語 Wikipedia で学習した SentencePiece のモデルで語彙数は 32000 です（SentencePiece のモデル作成方法は第7回を参考にして下さい）。

from t5.data import utils
utils.DEFAULT_SPM_PATH = "gs://somewhere/t5/sentencepiece/wikipedia_20190301_ja_v003.model"

以下のようにして日本語 Wikipedia を使う学習タスクを追加します。

import functools
from t5.data import preprocessors

from t5.data.dataset_providers import TaskRegistry
from t5.seqio import MixtureRegistry
from t5.data.dataset_providers import TfdsTask
from t5.data.tasks import DEFAULT_OUTPUT_FEATURES

task_name_wikipedia_ja = "wikipedia_20190301.ja_v003_unsupervised"

TaskRegistry.add(
    task_name_wikipedia_ja,
    TfdsTask,
    tfds_name="wikipedia/20190301.ja:1.0.0",
    text_preprocessor=functools.partial(
        preprocessors.rekey, key_map={"inputs": None, "targets": "text"}),
    token_preprocessor=preprocessors.unsupervised,
    output_features=DEFAULT_OUTPUT_FEATURES,
    metric_fns=[])

MixtureRegistry.add(task_name_wikipedia_ja, [(task_name_wikipedia_ja, 1.0)])

TPU のアドレスを確認して、

import os
import pprint
import json
import tensorflow.compat.v1 as tf

assert 'COLAB_TPU_ADDR' in os.environ, 'ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!'
TPU_ADDRESS = 'grpc://' + os.environ['COLAB_TPU_ADDR']
print('TPU address is', TPU_ADDRESS)

以下のようにしてハイパーパラメータを設定します。

from t5.models.mesh_transformer_main import FLAGS
from t5.models.mesh_transformer_main import main

FLAGS.mark_as_parsed()

FLAGS.tpu = TPU_ADDRESS
FLAGS.model_dir = 'gs://somewhere/t5/wikipedia_20190301.ja_v003/model'
tf.flags.FLAGS.gin_file=[
  "dataset.gin", 
  "models/bi_v1.gin",
  "objectives/span_3_15_u_u.gin",
  "learning_rate_schedules/rsqrt_no_ramp_down.gin"]
tf.flags.FLAGS.gin_param=[
  "utils.tpu_mesh_shape.model_parallelism = 1",
  "utils.tpu_mesh_shape.tpu_topology = 'v2-8'",
  "run.batch_size = ('tokens_per_batch', 65536)",
  "run.train_steps = 524288",
  "MIXTURE_NAME = 'wikipedia_20190301.ja_v003_unsupervised'"
]

補足:T5 のバージョン

ここで少しだけ脱線しますが、現在 T5 には 1.0, 1.1 の二つのバージョンがあります。 上記のコードでは第7回で使用したのと同じ "models/bi_v1.gin" を使用していますが、これが T5 の 1.0 です。 "models/t5.1.1.base.gin" と指定すると 1.1 になります。

1.0 からの変更点は以下のとおりです。詳しくはこちら⁸を参照して下さい。

• feedforward 層のアクティベーションを RELU から GEGLU に変更。
• 上記に伴い feedforward 層のユニット数を 2/3 に削減。
• 事前学習時の dropout を廃止(ただしファインチューニング時は使用する)
• 語彙数で分類する際の重みをトークンIDから埋め込み表現にする際の重みと共有しない。

最後に以下のようにして実行です。

tf.disable_v2_behavior()
tf.logging.set_verbosity(tf.logging.INFO)
main([])

当然ですが事前学習は Colab のランタイムのライフスパンでは終わりません。 動作確認した後は GCP で VM と Cloud-TPU を使って学習した方が手っ取り早いでしょう。

3. Transformers でのファインチューニング

ここからは先ほどの事前学習済みモデルを Transformers の形式に変換してファインチューニングをしていきます。 それでは新たにノートブックを開き、アクセラレータは GPU を選んで下さい。

セットアップ

事前学習で生成したチェックポイントを読み込むので GCS の認証を通します。

from google.colab import auth
auth.authenticate_user()

Transformers と SentencePiece をインストールします。

!pip install transformers==4.3.3
!pip install sentencepiece

datasets は Hugging Face が公開する自然言語処理のデータセットとメトリクスを利用しやすくするライブラリです。 自前のデータを読み込んで使うこともできるので、インストールしておきます。

!pip install datasets

ファインチューニングのメトリクスには BLEU を使うので、sacrebleu をインストールします。こちらは日本語に対応するようになりました。

!pip install sacrebleu[ja]

事前学習済みモデルの変換

まずは、t5 のチェックポイントを Transformers 向けに変換します。 チェックポイントと SentencePiece のモデルを手元のディレクトリに集めて、

!mkdir pretrained
!gsutil cp gs://somewhere/t5/wikipedia_20190301.ja_v003/model/*  ./pretrained
!gsutil cp gs://somewhere/t5/sentencepiece/wikipedia_20190301_ja_v003.*  ./pretrained

中身はこんな感じです。

!ls pretrained
# checkpoint model.ckpt-524288.meta
# model.ckpt-524288.data-00000-of-00002  operative_config.gin
# model.ckpt-524288.data-00001-of-00002  wikipedia_20190301_ja_v003.model
# model.ckpt-524288.index wikipedia_20190301_ja_v003.vocab

変換後のモデルの出力先のディレクトリを作って、

!mkdir t5_base_wikipedia_ja

必要なクラスをインポートです。

from transformers import T5Tokenizer, T5Config, T5Model

まずはトークナイザを保存します。

tokenizer = T5Tokenizer("./pretrained/wikipedia_20190301_ja_v003.model")
tokenizer.save_pretrained("./t5_base_wikipedia_ja")
# ('./t5_base_wikipedia_ja/tokenizer_config.json',
# './t5_base_wikipedia_ja/special_tokens_map.json',
# './t5_base_wikipedia_ja/spiece.model',
# './t5_base_wikipedia_ja/added_tokens.json')

ここからは T5 の変換です。まずは T5Config を作って JSON に保存します。

config = T5Config(
    decoder_start_token_id=0,
    vocab_size=32128, 
    n_positions=512, 
    d_model=768, 
    d_kv=64, 
    d_ff=3072, 
    num_beams=4,
    num_layers=12, 
    num_heads=12, 
    relative_attention_num_buckets=32, 
    dropout_rate=0.1, 
    layer_norm_epsilon=1e-06, 
    initializer_factor=1.0, 
    is_encoder_decoder=True, 
    pad_token_id=0, 
    eos_token_id=1)
config.to_json_file("./config.json")

以下のようにしてチェックポイントを変換します。

from transformers.models.t5.convert_t5_original_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch
convert_tf_checkpoint_to_pytorch(
    tf_checkpoint_path="./pretrained/model.ckpt-524288", 
    config_file="./config.json", 
    pytorch_dump_path="./pytorch_model.bin")
# ...
# Weights not copied to PyTorch model: 
# Configuration saved in ./pytorch_model.bin/config.json
# Save PyTorch model to ./pytorch_model.bin
# Model weights saved in ./pytorch_model.bin/pytorch_model.bin

最後の Weights not copied to PyTorch model: のところで出力がなければ、全てのパラメータが読み込まれています。 一旦、読み込んで所定のディレクトリに保存します。

model = T5Model.from_pretrained('./pytorch_model.bin')
model.save_pretrained("./t5_base_wikipedia_ja")
# Configuration saved in ./t5_base_wikipedia_ja/config.json
# Model weights saved in ./t5_base_wikipedia_ja/pytorch_model.bin

!ls ./t5_base_wikipedia_ja
# config.json special_tokens_map.json  tokenizer_config.json
# pytorch_model.bin  spiece.model

これでめでたく変換ができました。

学習データの準備

ファインチューニング用のデータセットは第7回と同じくやさしい日本語データセットを使います。 変換の手順だけ記述しますので、詳しくは第7回を参照して下さい。

!apt-get install nkf
!pip install python-Levenshtein

import Levenshtein
def levenshtein_distance(row):
  return Levenshtein.distance(row["target"], row["input"])

!curl -L -o T15.xlsx https://filedn.com/lit4DCIlHwxfS1gj9zcYuDJ/SNOW/T15-2020.1.7.xlsx
!curl -L -o T23.xlsx https://filedn.com/lit4DCIlHwxfS1gj9zcYuDJ/SNOW/T23-2020.1.7.xlsx

import pandas as pd
snow_t15 = pd.read_excel('T15.xlsx')
snow_t15 = snow_t15.rename(columns={'#日本語(原文)': 'input', '#やさしい日本語':'target'})[['ID', 'input','target']]

snow_t15["input_len"] = snow_t15["input"].apply(len)
snow_t15["target_len"] = snow_t15["target"].apply(len)
snow_t15['levenshtein_distance'] = snow_t15.apply(levenshtein_distance, axis=1)

snow_t15 = snow_t15.query('levenshtein_distance < 10')[['input', 'target']]
snow_t15.to_csv("temp.tsv", sep="\t", header=False, index=False)
!cat temp.tsv | nkf -m0Z1 | tr "[:upper:]" "[:lower:]" > snow_t15.tsv

import pandas as pd
snow_t23 = pd.read_excel('T23.xlsx')
snow_t23 = snow_t23.rename(columns={'#日本語(原文)': 'input', '#やさしい日本語':'target'})[['ID', 'input','target']]

snow_t23["input_len"] = snow_t23["input"].apply(len)
snow_t23["target_len"] = snow_t23["target"].apply(len)
snow_t23['levenshtein_distance'] = snow_t23.apply(levenshtein_distance, axis=1)

snow_t23 = snow_t23.query('levenshtein_distance < 10')[['input', 'target']]
snow_t23.to_csv("temp.tsv", sep="\t", header=False, index=False)
!cat temp.tsv | nkf -m0Z1 | tr "[:upper:]" "[:lower:]" > snow_t23.tsv

!cat snow_t15.tsv snow_t23.tsv > snow.tsv

with open("snow.tsv", "r") as f:
  lines = f.readlines()
  lines = [line.strip() for line in lines]

from sklearn.model_selection import train_test_split
train, dev_test = train_test_split(lines, train_size=0.8, random_state=4)
dev, test = train_test_split(dev_test, train_size=0.5, random_state=7)

with open("snow_t15_23_train.tsv", "w") as f:
    f.write("\n".join(train)+"\n")

with open("snow_t15_23_dev.tsv", "w") as f:
    f.write("\n".join(dev)+"\n")

with open("snow_t15_23_test.tsv", "w") as f:
    f.write("\n".join(test)+"\n")

!wc -l snow_t15_23_*.tsv
#   7045 snow_t15_23_dev.tsv
#   7045 snow_t15_23_test.tsv
#  56357 snow_t15_23_train.tsv
#  70447 total

と思ったら動かすと2か所程何やらエラーがでます。学習データ(↓)の 42434 行目と、

!cat -n snow_t15_23_train.tsv | head -42436 | tail -5
# 42432 私の知っていることといえば、彼が中国からやってきたということだけです。   私の知っていることといえば、彼が中国からやってきたことだけです。
# 42433 彼はその劇でわき役を演じた。  彼はその劇で一番ではない役になった。
# 42434 "
# 42435 ブラウンさんの電話番号を調べてください。    ブラウンさんの電話番号を調べてください。
# 42436 彼女は1週間前に病気になった。 彼女は1週間前に病気になった。

テストデータ(↓)の 328 行目です。右列の先頭に地味に「"」が入ってますね。。。

cat -n snow_t15_23_test.tsv | head -330 | tail -5
# 326   もう何をするか決めましたか。  もう何をするか決めましたか。
# 327   さらに厄介なことに、彼は近所の人に迷惑をかけていることにさえ気づいていない。  さらに大変なことに、彼は家の近くの人に迷惑をかけていることにさえ気づいていない。
# 328   現在の彼女は結婚前の彼女ではない。 "今の彼女は結婚前の彼女ではない。
# 329   ここからヒルトンホテルまでどのくらい時間がかかりますか。    ここからヒルトンホテルまでどのくらい時間がかかりますか。
# 330   よく悪夢を見ます。 よく悪い夢を見ます。

真面目に対応するのも面倒くさかったので、行ごと消しちゃいました。

!sed -i '/^"/d' snow_t15_23_train.tsv
!sed -i -e '328d' snow_t15_23_test.tsv

では、先に進みましょう。

事前学習済みモデルのロード

それでは改めて先ほど変換した事前学習済みモデルをロードします。使用クラスが T5ForConditionalGeneration になっているので気を付けて下さいね。

from transformers import T5Tokenizer, T5Config, T5ForConditionalGeneration
tokenizer = T5Tokenizer.from_pretrained("./t5_base_wikipedia_ja")
model = T5ForConditionalGeneration.from_pretrained("./t5_base_wikipedia_ja")

ファインチューニングの実行

ここからファインチューニングを実行していきます。Transformers には Trainer という便利なクラスがあるので使ってみましょう。 Trainer は機械学習にありがちな学習ループ等を隠ぺいしてくれるクラスです。

まずは、ハイパーパラメータはこんな感じにしました。

MAX_SEQ_LEN = 64
TRAIN_BATCH_SIZE  = 8
EVAL_BATCH_SIZE  = 8
NUM_EPOCHS  = 3

使用するクラスのインポートです。

import torch
from datasets import load_dataset
from transformers import Trainer, TrainingArguments
from torch.utils.data import DataLoader

以下のようにしてデータセットを定義します。

dataset = load_dataset('csv', delimiter="\t", 
                       column_names = ["input", "output"],
                       data_files={"train": "./snow_t15_23_train.tsv", 
                                   "validation": "./snow_t15_23_dev.tsv",
                                   "test": "./snow_t15_23_test.tsv"})

こんな感じで中身を確認できます。

dataset["train"][0]
# {'input': '君はその計画を予定通り実行すべきだ。', 'output': 'あなたはその計画を予定通り実行するべきだ。'}

トークナイズとパディングの処理です。返却する dict のキーをどう定義したらよいのか謎だったのですが、モデルへの入力は forward() メソッド⁹に合わせて "input_ids" としておけば良いようです。モデルの出力と突き合わせるラベルの方は Trainer の compute_loss() メソッドを見ると¹⁰、"labels" にしておけば良さそうです。

def tokenize_function(examples):
    return {"input_ids": tokenizer(examples["input"], max_length=MAX_SEQ_LEN, padding="max_length").input_ids,
            "labels": tokenizer(examples["output"], max_length=MAX_SEQ_LEN, padding="max_length").input_ids}
tokenized_dataset = dataset.map(tokenize_function, batched=True)

これを PyTorch のテンソルにしておきます。

tokenized_dataset.set_format(type='torch', columns=['input_ids', 'labels'])

こちら¹¹の議論では、オプティマイザは Adafactor が良いようです。学習レートは T5 の論文に倣って 0.001 で固定としたそうで。 fairseq ¹² の実装を使ったようですが、学習パラメータの定義クラスである TrainingArguments を見ていると adafactor というフラグを見つけました。実装をチェックすると、

#cited from https://github.com/huggingface/transformers/blob/bae0c79f6fab7b1156dbb769c5493d1b393d7fa6/src/transformers/trainer.py#L578-#L588
578|            optimizer_cls = Adafactor if self.args.adafactor else AdamW
579|            if self.args.adafactor:
580|                optimizer_cls = Adafactor
581|                optimizer_kwargs = {"scale_parameter": False, "relative_step": False}
...|            ...
588|            optimizer_kwargs["lr"] = self.args.learning_rate

このフラグで良さそうですね。 学習レートを固定にする設定ですが、学習レートのスケジューリングの処理は以下のようになっています。

#cited from https://github.com/huggingface/transformers/blob/bae0c79f6fab7b1156dbb769c5493d1b393d7fa6/src/transformers/trainer.py#L598-#L604
598|       if self.lr_scheduler is None:
599|            self.lr_scheduler = get_scheduler(
600|                self.args.lr_scheduler_type,
601|                self.optimizer,
602|                num_warmup_steps=self.args.warmup_steps,
603|                num_training_steps=num_training_steps,
604|            )

get_scheduler() の引数になっている SchedulerType の定義を確認しておきましょう。

#cited from https://github.com/huggingface/transformers/blob/bae0c79f6fab7b1156dbb769c5493d1b393d7fa6/src/transformers/trainer_utils.py#L226-L232
226| class SchedulerType(ExplicitEnum):
227|     LINEAR = "linear"
228|     COSINE = "cosine"
229|     COSINE_WITH_RESTARTS = "cosine_with_restarts"
230|     POLYNOMIAL = "polynomial"
231|     CONSTANT = "constant"
232|     CONSTANT_WITH_WARMUP = "constant_with_warmup"

これなら lr_scheduler_type に "constant" を設定しておけばよさそうですね。

そんな訳でTrainer のパラメータを以下のように設定しました。 オプティマイザに学習レート 0.001 固定の AdaFactor を使い、2000 ステップ毎に checkpoint を保存して、同じタイミングでログ出力と検証を行います。

training_args = TrainingArguments(
    "./finetune",
    num_train_epochs = NUM_EPOCHS,
    evaluation_strategy = "steps",
    adafactor=True,
    learning_rate=1e-3,
    lr_scheduler_type="constant",
    per_device_train_batch_size = TRAIN_BATCH_SIZE,
    per_device_eval_batch_size  = EVAL_BATCH_SIZE,
    eval_steps = 2000,
    logging_steps = 2000,
    save_steps = 2000,
)
)

GPU が使用されるのを確認しておきます。

training_args.device
# device(type='cuda', index=0)

Trainer を定義します。compute_metrics でメトリクス関数を設定できるはずなのですが、BLEU の計算をしこんだら OOM になってしまったので諦めました。。。

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["validation"],
    #compute_metrics = compute_metrics
)

以下のようにして学習ループを実行します。

trainer.train()

14000 ステップ辺りがよさそうです。チェックポイントからパラメータを読み込み、名前を付けて保存しておきます。

model = T5ForConditionalGeneration.from_pretrained("./finetune/checkpoint-14000")
model.save_pretrained("./t5_snow")

学習を回した後で知ったのですが、TrainingArguments に load_best_model_at_end = True を設定しておくと、学習後に最良のモデルを自動でロードしてくれるらしいです。

それではテストデータを使って推論してみましょう。

テストデータでの推論

テストデータを読込みます。

test_dl = DataLoader(tokenized_dataset["test"], shuffle=False, batch_size=EVAL_BATCH_SIZE)

GPU を使えるなら GPU を使いますという設定です。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

先ほど保存したファインチューニング済みのモデルをロードして、こんな感じで動かします。

model = T5ForConditionalGeneration.from_pretrained("./t5_snow")
model.to(device)
model.eval()
predictions = []
for i, batch in enumerate(test_dl):
  if i % 100 == 0:
    print("step = %d" % i)
  inputs = batch["input_ids"].to(device)
  outputs = model.generate(inputs, max_length=MAX_SEQ_LEN, num_beams=4)
  predictions.extend(outputs.cpu().numpy())
# step = 0
# step = 100
# ...
# step = 700
# step = 800

入力した文章はこんな感じで、

[example["input"] for example in dataset["test"]][:10]
#['ほとんど何も知りません。',
# '彼女は彼に愛されている。',
# '私は言葉で気持ちを伝えることができます。',
# '君は何を言おうとしているの。',
# '彼は一週間学校を休んだ。',
# '私は妻を見舞いに病院にいった。',
# '傘を持ってきましたか。',
# '彼には審美眼がない。',
# '彼の家族は早起きです。',
# '誤解を正させてください。']

お手本の参照文はこんな感じ、

references = [example["output"] for example in dataset["test"]]
references[:10]
#['ほとんど何も知りません。',
# '彼女は彼に愛してもらっている。',
# '私は言葉で気持ちを伝えることができます。',
# 'あなたは何を言おうとしているの。',
# '彼は1週間学校を休んだ。',
# '私は妻を見に病院にいった。',
# '傘を持ってきましたか。',
# '彼には美しさを判断する力がない。',
# '彼の家族は早く起きます。',
# '間違った理解を直させてください。']

生成したテキストはこんな感じです。まぁそれなりにそれっぽいでしょうか。

decoded_texts = [tokenizer.decode(prediction, skip_special_tokens=True) for prediction in predictions]
decoded_texts[:10]
#['ほとんど何も知りません。',
# '彼女は彼に愛してもらっている。',
# '私は言葉で気持ちを伝えることができます。',
# 'あなたは何を言おうとしているの。',
# '彼は1週間学校を休んだ。',
# '私は妻の様子を見に病院にいった。',
# '傘を持ってきましたか。',
# '彼には警察の目がない。',
# '彼の家族は早く起きます。',
# '間違った理解を正させてください。']

それっぽいで終わるのもなんなので、 BLEU を計算してみましょう。

from sacrebleu import corpus_bleu

def bleu(predictions, references):
  references = [references]
  bleu_score = corpus_bleu(predictions, references,      
                                     smooth_method="exp",
                                     smooth_value=0.0,
                                     force=False,
                                     lowercase=False,
                                     tokenize="ja-mecab",
                                     use_effective_order=False)
  return bleu_score.score

bleu(decoded_texts, references)
# 78.5876703725097

第7回で Google の実装を使ったときは、BLEU = 78.795 だったのでそこそこ近い値になりました （第7回のときと、sacrebleu のバージョンが違うとか、日本語で BLEU の計算するために怪しげにいじくってたとかありますが）。 以前に Transformers で T5 がサポートされた直後に Adam で試したときは、72.42 だったのでだいぶ改善しました。

次は FP16 を試してみましょう。

4. FP16 モードでの実行

通常はモデルのパラメータは単精度浮動小数点(FP32)で動作していますが、 Transformers では非常に簡単に半精度浮動小数点(FP16)にパラメータの型を変換することができます。「32bit の数値を 16bit にしたら情報が落ちちゃうんじゃないか？大丈夫なの？」という気がしますが、意外と大丈夫らしく、モデルやタスクにもよりますが、 FP32 で学習したモデルのパラメータを FP16 に変換して推論してもそれ程精度は変わらなかったりします。

ただ、 T5 に関してはそんな簡単な話ではなかったようで、まがりなりにも FP16 で動くようになったのは今年に入ってからでした。 なん箇所かクリッピングの処理を入れたみたいですね。

どうして、FP16 で動かしたいかというと、サイズ半分だから同じ量の GPU メモリで余分にモデルをデプロイできたり、推論速度が速くなったりというところです。

それでは T5 を FP16 で動かして見ましょう。

FP16 で動かすのは簡単で、model.half() するだけでOKです。

model.half()

実際にパラメータの型を覗いてみましょう。ちゃんと FP16 になっています。

model.shared.weight.dtype
# torch.float16

推論もちゃんと動きます。

model.to(device)
model.eval()
input_ids = tokenizer.encode("人生は空虚な夢ではない。", max_length=64, pad_to_max_length=True, return_tensors="pt")  
inputs = input_ids.to(device)
outputs = model.generate(inputs, max_length=64, num_beams=4)
decoded_text = tokenizer.decode(outputs.cpu().numpy()[0], skip_special_tokens=True)
decoded_text
# 人生は悲しい夢ではない。

精度にどの程度影響があるか、もう一度 BLEU を計算してみましょう。

predictions = []
for i, batch in enumerate(test_dl):
  if i % 100 == 0:
    print("step = %d" % i)
  inputs = batch["input_ids"].to(device)
  outputs = model.generate(inputs, max_length=MAX_SEQ_LEN, num_beams=4)
  predictions.extend(outputs.cpu().numpy())

# step = 0
# step = 100
# ...
# step = 700
# step = 800

decoded_texts = [tokenizer.decode(prediction, skip_special_tokens=True) for prediction in predictions]
bleu(decoded_texts, references)
# 78.58615166254424

タスクの難易度にもよるのでしょうが、今回のようなシンプルな書き換え問題だと殆ど変わらないですね。 それでは、実際に速くなるの？というところを見ていきましょう。

5. Tensor Core

モデルを FP16 で高速に推論する場合 Tensor Core について考慮する必要があります。 Tensor Core は NVIDIA の Volta アーキテクチャ以降の GPU に搭載されている計算実行ユニットで、4 x 4 の行列 A, B, C があったとして、 AB+C を 1 サイクルで計算することができます。Tensor Core を使った計算でサポートされるデータ型はアーキテクチャにより異なりますが FP16 は Volta 以降の全てのアーキテクチャで利用可能です。

また、どんなモデルでも Tensor Core で動かせる訳ではなくて、制限もあります。こちらの資料¹³の 23 ページにまとまっていたので引用しておきます。

ざっくり次元数は 8 の倍数、できたら 64 の倍数がよいということでしょうか。 事前学習済みのモデルがあるので次元数とか今さら変えようないのですが、改めて今回のモデルの構造を見ておきましょう。

print(model)

# T5ForConditionalGeneration(
#   (shared): Embedding(32128, 768)
#   (encoder): T5Stack(
#     (embed_tokens): Embedding(32128, 768)
#     (block): ModuleList(
#       (0): T5Block(
#         (layer): ModuleList(
#           (0): T5LayerSelfAttention(
#             (SelfAttention): T5Attention(
#               (q): Linear(in_features=768, out_features=768, bias=False)
#               (k): Linear(in_features=768, out_features=768, bias=False)
#               (v): Linear(in_features=768, out_features=768, bias=False)
#               (o): Linear(in_features=768, out_features=768, bias=False)
#               (relative_attention_bias): Embedding(32, 12)
#             )
#             (layer_norm): T5LayerNorm()
#             (dropout): Dropout(p=0.1, inplace=False)
#           )
#           (1): T5LayerFF(
#             (DenseReluDense): T5DenseReluDense(
#               (wi): Linear(in_features=768, out_features=3072, bias=False)
#               (wo): Linear(in_features=3072, out_features=768, bias=False)
#               (dropout): Dropout(p=0.1, inplace=False)
#             )
#             (layer_norm): T5LayerNorm()
#             (dropout): Dropout(p=0.1, inplace=False)
#           )
#         )
#       )
#       ...
#  
#       (11): T5Block(
#         (layer): ModuleList(
#           (0): T5LayerSelfAttention(
#             (SelfAttention): T5Attention(
#               (q): Linear(in_features=768, out_features=768, bias=False)
#               (k): Linear(in_features=768, out_features=768, bias=False)
#               (v): Linear(in_features=768, out_features=768, bias=False)
#               (o): Linear(in_features=768, out_features=768, bias=False)
#             )
#             (layer_norm): T5LayerNorm()
#             (dropout): Dropout(p=0.1, inplace=False)
#           )
#           (1): T5LayerCrossAttention(
#             (EncDecAttention): T5Attention(
#               (q): Linear(in_features=768, out_features=768, bias=False)
#               (k): Linear(in_features=768, out_features=768, bias=False)
#               (v): Linear(in_features=768, out_features=768, bias=False)
#               (o): Linear(in_features=768, out_features=768, bias=False)
#             )
#             (layer_norm): T5LayerNorm()
#             (dropout): Dropout(p=0.1, inplace=False)
#           )
#           (2): T5LayerFF(
#             (DenseReluDense): T5DenseReluDense(
#               (wi): Linear(in_features=768, out_features=3072, bias=False)
#               (wo): Linear(in_features=3072, out_features=768, bias=False)
#               (dropout): Dropout(p=0.1, inplace=False)
#             )
#             (layer_norm): T5LayerNorm()
#             (dropout): Dropout(p=0.1, inplace=False)
#           )
#        )
#       )
#     )
#     (final_layer_norm): T5LayerNorm()
#     (dropout): Dropout(p=0.1, inplace=False)
#   )
#   (lm_head): Linear(in_features=768, out_features=32128, bias=False)
# )     

Linear が沢山ありますが入力と出力の次元数は 768 や 3072 や 32128 で、とりあえず 8 の倍数ではありますね。 batch 軸に関しては“入力文"のバッチサイズ × シーケンス長になるんではないかと思います。なのでシーケンス長を 8 の倍数にしてパディングすればいいのかなと。

そんな訳で FP16 化したモデルを Tensor Core で動かして、どの程度速くなるのか確認して見ました。 Tensor Core が必要なので環境としては AWS の g4xn.xlarge で Tesla T4 (Turing アーキテクチャ) で実験をしています。

以降、入力シーケンス長を 192(もしくはパディングなし)、出力シーケンス長を 64 にしていますが深い意味はありません。 FP32 に対して FP16 で実行速度が上がるかどうかが興味の対象だったので。

実行環境

実行環境は以下のとおりです。

• Instance Type : g4dn.xlarge (Tesla T4)
• Driver version : 460.32.03
• Docker Image : nvcr.io/nvidia/pytorch:20.12-py3 (CUDA 11.1.1)

実験コード

実験に使用したコード(test_t5_tensorcore.py)は以下のとおりです。モデルをロードした後、ウォームアップとして 1 回推論、その後に 10 回推論するのに要した時間を time.time() で計測しています。

# -*- coding: utf-8 -*-
import argparse
import os
import time
import sys
import json
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration

# $ docker run --gpus all --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 -it nvcr.io/nvidia/pytorch:20.12-py3 bash
# # pip install transformers==4.2.1
# # nvprof test_t5_tensorcore.py > nvprof.txt 2>&1

parser = argparse.ArgumentParser()

parser.add_argument('--fp16', action='store_true')
parser.add_argument('--batch_size', type=int, default=8)
parser.add_argument('--input_seq_len', type=int, default=None)
parser.add_argument('--output_seq_len', type=int, default=64)
parser.add_argument('--beam_size', type=int, default=4)

args = parser.parse_args()

t5_model_path = "./t5_base_finetuned"
cuda_device = "cuda:0"

print("fp16           = %s" % (args.fp16))
print("batch_size     = %d" % (args.batch_size))
print("beam_size      = %d" % (args.beam_size))
print("input_seq_len  = %s" % (args.input_seq_len))
print("output_seq_len = %s" % (args.output_seq_len))

use_fp16 = args.fp16 
batch_size = args.batch_size 
input_seq_len = args.input_seq_len
output_seq_len = args.output_seq_len
num_beams = args.beam_size

tokenizer = T5Tokenizer.from_pretrained(t5_model_path)
model = T5ForConditionalGeneration.from_pretrained(t5_model_path)

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

model.eval()
if use_fp16:
  model.half()

inputs = ["おはようございます。"] * batch_size

if input_seq_len :
  input_ids = tokenizer(inputs, max_length=input_seq_len, padding='max_length', return_tensors="pt").input_ids.to(device) 
else:
  input_ids = tokenizer(inputs, return_tensors="pt").input_ids.to(device) 
print("shape of input_ids : %s" % (input_ids.shape,))

output_ids = model.generate(input_ids, max_length=output_seq_len, num_beams=num_beams)

start = time.time()
for i in range(10):
  output_ids = model.generate(input_ids, max_length=output_seq_len, num_beams=num_beams)
process_time = time.time() - start
output = tokenizer.decode(output_ids.cpu().numpy()[0], skip_special_tokens=True)
print("input: %s, output: %s" % (inputs[0], output))

print("Elapsed time : %s" % (process_time))

Tensor Core 利用の確認

上記のコードを動かせば、とりあえず実行時間を計測できますが、 Tensor Core が本当に使われてるかどうか気になります。 簡易的な確認の仕方はさっきの資料¹³の 10 ページに記載があり、"Run nvprof and look for [i|s|h][some numbers] in function names" だそうです。

上記のコードであれば以下のように動かして、ログを確認します。

 nvprof python test_t5_tensorcore.py --batch_size 1 --input_seq_len 192 --fp16 > nvprof_fp16_in_192_bs_1.log 2>&1

設定を変えて何パターンか試すと以下のようになりました(設定したパラメータはファイル名から察して下さい。。。)。

!grep -cHe "\_[ish][0-9+]" nvprof*.log
# nvprof_fp16_in_192_bs_1.log:6
# nvprof_fp16_in_192_bs8.log:9
# nvprof_fp16_in_none_bs_1.log:2
# nvprof_fp16_in_none_bs_8.log:4
# nvprof_fp32_in_192_bs_1.log:0
# nvprof_fp32_in_192_bs_8.log:0
# nvprof_fp32_in_none_bs_1.log:0
# nvprof_fp32_in_none_bs_8.log:0

ログから読み取れたこととして、

• 当然ながら、 FP32 を指定した場合は Tensor Core が使用された形跡はありません。
• 入力長を 192 に揃えるかどうか(in_192 と in_none の比較) では条件に引っかかった行数が異なりますが、in_none の方は入力長が 8 の倍数じゃないので、encoder 側の計算で Tensor Core が使用されなかったものと思われます。
• バッチサイズが 1 の時でも Tensor Core が使えていますが、これは内部的には Linear のバッチ軸が バッチサイズ×シーケンス長になっているからではないかと思います。

バッチサイズを 8 の倍数にしてあげた方が具合が良さそうにも見えますが、バッチを組むとデコーダの計算が最も出力長が長くなるサンプルに引っ張られそうなので悩ましいですね。

何はともあれ Tensor Core は使用されているようです。後はどれだけ速くなったか確認してみましょう。

実行速度の確認

肝心の実行速度です。

!grep -H "Elapsed" nvprof*.log
# nvprof_fp16_in_192_bs_1.log:Elapsed time : 7.19362211227417
# nvprof_fp16_in_192_bs_8.log:Elapsed time : 8.30354380607605
# nvprof_fp16_in_none_bs_1.log:Elapsed time : 7.086625814437866
# nvprof_fp16_in_none_bs_8.log:Elapsed time : 8.016315460205078
# nvprof_fp32_in_192_bs_1.log:Elapsed time : 7.60028338432312
# nvprof_fp32_in_192_bs_8.log:Elapsed time : 10.066397428512573
# nvprof_fp32_in_none_bs_1.log:Elapsed time : 7.443647861480713
# nvprof_fp32_in_none_bs_8.log:Elapsed time : 8.863409757614136

あれ、パッと見た感じ速くなってる気がしませんね。。。もうちょっと整理するとこんな感じです。

Tensor Core 使った具合からいうと、 2 行目のバッチサイズ = 8, 入力シーケンス長 = 192 が速度差が出て欲しいところですが、1.2 倍でした。 速くなったかと言われれば速くなっていますが、何倍! みたいなのを期待していたので、がっかりですね。

本当はここで「○倍高速化しましたね！では REST API にしてみましょう!」と続く予定だったのですけどね。。。 ややテンション下がり気味ですが、 REST API にしてアプリから呼べるようにしてみましょう。

6. REST API 化

さて、速くならないのは残念ですが、気を取り直して REST API にしてみましょう。

TorchScript化？

Tensorflow の場合は SavedModel にエクスポートして、Tensorflow Serving にデプロイすれば OK だったのですが、 Transformers で PyTorch の場合、torch.jit.trace() で TorchScript ¹⁴ 化して、それを Triton ¹⁵にデプロイするのが良さそうです。

まず T5 のドキュメント¹⁶を読むと、テキスト生成する場合は generate() メソッドの使用が推奨されています。

• For sequence-to-sequence generation, it is recommended to use T5ForConditionalGeneration.generate().

次に T5 の generate() をどうやって TorchScript 化するか調べ始めて、この issue ¹⁷ に突き当たりました。 Hugging Face の中の人が以下のようにおっしゃっています。。。

• Yeah, I don’t think our generate() method is torchscriptable yet :-/

確かに issue に記載されているコード¹⁸で TorchScript 化はできましたけど、これを Triton にデプロイすると 「入力トークン列と出力トークン列を入力して次の出力トークンを得る」ような API となり、呼び出し元がぐるぐる呼び出しながら自前でビーム探索なりなんなり実装することになりそうです。。。

面倒な上に遅そうなので別な作戦を考えることになりました。

Gunicorn/Falcon による実装

どうしたものか？と考えて検索したところ、別の Hugging Face の中の人の記事¹⁹ がヒットです。 ざっと目を通すと以下のような話です。細かい相違点はありますが、大まかには今回の T5 でテキスト生成をする局面にマッチしてそうです。

• GPT-2 でテキスト生成をする
• バッチ処理ができない局面を想定
• 複数リクエストを並行処理して GPU を効率良く使いたい

実装まで見つからなかったのですが、基本的には Falcon で実装したシングルスレッドの API をワーカとし、それを Gunicorn で複数ワーカを並行稼働させる作りのようです。構造的にはちょうど以下のようにすれば良いかと思います。

簡易的にですが見よう見まねで作ってみたので、ご紹介しておきましょう。

ディレクトリ構造とファイルはこんな感じです。

/somewhere/
 + docker/
   + Doclerfile
   + t5mode.py
 + conf/
   + gunicorn.conf.py
   + t5conf.json
   + t5_base_finetuned/

各ファイルの内容を見ていきましょう。

Dockerfile

NVIDIA の PyTorch イメージをベースに必要なものをインストールしたり、コピーしたりして、Gunicorn の起動設定をしているだけですね。 ちなみに NVIDIA のイメージを使っているのは nvprof が入っていたほうが便利かな？と思ったからというだけです。

FROM nvcr.io/nvidia/pytorch:20.12-py3
RUN pip install falcon==2.0.0 gunicorn==20.0.4 transformers==4.2.1
RUN mkdir /server
WORKDIR /server
ADD ./t5model.py /server

ENTRYPOINT ["/opt/conda/bin/gunicorn", "t5model"]
CMD ["-c", "/conf/gunicorn.conf.py"]

t5model.py

メインになる Falcon のロジックです。正直 logger の設定とかは見たい情報出たらいいやレベルでテキトーに、エラー処理とかも省略してます。

import falcon
import os
import sys
import json
import logging
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration

if "T5CONF" in os.environ:
    conf_path = os.environ["T5CONF"]
else:
    conf_path = None

T5CONF = {
  't5_model_path': './t5_base_wikipedia_ja',
  'cuda_device': 'cuda:0',
  'use_fp16': True,
  'max_input_length': None,
  'max_output_length': 64,
  'num_beams': 4,
  'log_level': "INFO"}

if conf_path and os.path.exists(conf_path):
  with open(conf_path, "r") as f:
    T5CONF.update(json.load(f))

root = logging.getLogger()
root.setLevel(T5CONF["log_level"])

logger = logging.getLogger(__name__)
formatter = logging.Formatter("[%(asctime)s] [%(process)d] [%(levelname)s] [%(name)s] %(message)s", datefmt='%Y-%m-%d %I:%M:%S %z')
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(T5CONF["log_level"])
handler.setFormatter(formatter)
logger.addHandler(handler)

class RequireJSON(object):
  def process_request(self, req, resp):
    if not req.client_accepts_json:
      raise falcon.HTTPNotAcceptable('Only supports responses encoded as JSON.')

    if req.method in ('POST', 'PUT'):
      if 'application/json' not in req.content_type:
        raise falcon.HTTPUnsupportedMediaType('Only supports requests encoded as JSON.')

class JSONTranslator(object):

  def process_request(self, req, resp):
    if req.content_length in (None, 0):
      return

    body = req.stream.read()
    if not body:
      raise falcon.HTTPBadRequest('A valid JSON document is required.')

    try:
      req.context.doc = json.loads(body.decode('utf-8'))
    except (ValueError, UnicodeDecodeError):
      raise falcon.HTTPError(falcon.HTTP_753, 'JSON was incorrect or not encoded as UTF-8.')

  def process_response(self, req, resp, resource, req_succeeded):    
    if not hasattr(resp.context, 'result'):
      return

    resp.body = json.dumps(resp.context.result)

class T5ModelResource(object):

  def __init__(self):
    self.t5_model_path = T5CONF["t5_model_path"]
    self.cuda_device = T5CONF["cuda_device"]
    self.use_fp16 = T5CONF["use_fp16"]
    self.max_input_length = T5CONF["max_input_length"] 
    self.max_output_length = T5CONF["max_output_length"] 
    self.num_beams = T5CONF["num_beams"]

    self.tokenizer = T5Tokenizer.from_pretrained(self.t5_model_path)
    self.model = T5ForConditionalGeneration.from_pretrained(self.t5_model_path)
    logger.info("T5 model is loaded from %s" % (self.t5_model_path))

    self.device = torch.device(self.cuda_device if torch.cuda.is_available() else "cpu")
    self.model.to(self.device)
    logger.info("Model parameters are placed on device %s" % (self.device))

    self.model.eval()
    if self.use_fp16:
      logger.info("Model parameters are converted to float16.")
      self.model.half()

    logger.info("Max input sequence length = %s." % (self.max_input_length))
    logger.info("Max output sequence length = %s." % (self.max_output_length))
    logger.info("Beam search with = %d" % (self.num_beams))

  def on_post(self, req, resp):
    logger.debug("requested json: %s" % (req.context.doc))
    input = req.context.doc["input"]
    if self.max_input_length:
      logger.debug("input_ids is padded to max_input_length = %d" % (self.max_input_length))
      input_ids = self.tokenizer(input, padding='max_length', 
                    max_length=self.max_input_length, return_tensors="pt").input_ids.to(self.device) 
    else:
      input_ids = self.tokenizer.encode(input, return_tensors="pt").to(self.device) 
    logger.debug("input shape : %s" % (input_ids.shape,))
    output_ids = self.model.generate(input_ids, 
      max_length=self.max_output_length, num_beams=self.num_beams)
    output = self.tokenizer.decode(output_ids.cpu().numpy()[0], skip_special_tokens=True)
    logger.info("input: %s, output: %s" % (input, output))
    output = {"output": output}
    resp.context.result = output
    resp.status = falcon.HTTP_200

application = falcon.API(middleware=[RequireJSON(), JSONTranslator(),])
t5 = T5ModelResource()
application.add_route('/t5', t5)

gunicorn.conf.py

Gunicorn の設定ファイルです。workers の値でワーカ数を変更します。デフォルトではワーカーはシングルスレッドで動くので、 workers で指定した数が並行処理を受け入れるリクエスト数になります。

workers = 1
bind = "0.0.0.0:8000"

t5.conf.json

T5 を初期化する設定ファイルです。

{
  "t5_model_path": "/conf/t5_base_finetuned",
  "cuda_device": "cuda:0",
  "use_fp16": true,
  "max_input_length": 192,
  "max_output_length": 64,
  "num_beams": 4,
  "log_level": "INFO"
}

t5_base_finetuned

ファインチューニング済みのディレクトリの中身はこうなってます。

$ ls conf/t5_base_finetuned
# config.json  special_tokens_map.json  tokenizer_config.json
# pytorch_model.bin  spiece.model

起動方法とテスト実行

docker ディレクトリで以下のようにしてコンテナのイメージをビルドします。

$ docker build -t t5-server .

コンテナのイメージができたら、以下のようにして起動します。

$ docker run --name t5-server --gpus all \
     --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 \
     -p 8080:8000 -e T5CONF=/conf/t5conf.json -v /somewhere/conf:/conf t5-server -c /conf/gunicorn.conf.py

実行するとこんな感じです。めでたく REST API 化できました。

curl -X POST -H "Content-Type: application/json" -d '{"input":"まあ少しの間はそれで大丈夫だろう。"}' localhost:8080/t5
# {"output": "まあ当分はそれで間に合うだろう。"}

これを g4dn.xlarge (Tesla T4) で起動して負荷をかけてみたところ、workers=4 の設定で以下のように GPU を使いきれてました。

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
| N/A   60C    P0    55W /  70W |   7087MiB / 15109MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      6324      C   /opt/conda/bin/python            1769MiB |
|    0   N/A  N/A      6325      C   /opt/conda/bin/python            1773MiB |
|    0   N/A  N/A      6326      C   /opt/conda/bin/python            1769MiB |
|    0   N/A  N/A      6330      C   /opt/conda/bin/python            1773MiB |
+-----------------------------------------------------------------------------+

スループットはモデルや投入するデータにもよりますが、

• 入力シーケンス長=192
• 出力シーケンス長=64
• ビーム幅=4 の設定で、

の条件設定で、

• 今回の Gunicorn/Falcon の実装
• Google 実装で生成した SavedModel を Tensorflow Serving にデプロイしたもの²⁰
• 同じ SavedModel を Triton にデプロイしたもの²¹

の3種類で比較したところ、以下のようになりました。

これからは T5 も Transformers を使えば OK ですね！というつもりが t5 のコードでエクスポートした SavedModel on Triton が最速になってしまいました。。。 FP16 でガッと速くなっていればよかったんですが。。。

7. おわりに

今回は Transformers の PyTorch 実装を使って Google 実装で生成した事前学習済みモデルの変換とファインチューニング、FP16化、REST API としての公開方法といったところをご紹介しました。

さて次回はどうしましょうか。。。。 spaCy が進化しているのでその話でもしようかな。