Google ColabのTPUでtransformerを学習

NLP 深層学習 自然言語処理 transformer CoLab TPU メモ

注: この記事は2019年4月29日現在のColabとTensorflow(1.13)での話です。

概要

kerasで書かれたtransformerをtf.kerasで書き直してGoogle Colabの無料で使えるTPU上で学習させた。

デモとして「Wikipedia日英京都関連文書対訳コーパス」を使って英→日翻訳を学習。
(入力・出力それぞれ)1024トークンx8を1バッチとしたものを8TPUで処理できることを確認。 また、サンプル(50万対弱)を5周するのにかかった時間は6483秒だった。

使用したリポジトリ:
https://github.com/iki-taichi/tf-keras-transformer

ノートブック:
https://github.com/iki-taichi/tf-keras-transformer/blob/master/notebook/tf_keras_transformer.ipynb

(ColabでPython3を開いてTPUを有効にし、ノートブックを再現可能)

今回リポジトリの使い方と、実装の変更点のメモ。

リポジトリを使ってColab TPUで学習するステップ

Colabについて参考にさせていただいた記事

[x] 今回、再現にはGCPのバケットは不要。

1. Google Colaboratory にアクセスしPython3ノートブックを開く

Colabにブラウザでアクセスして新しいPython3ノートブックを開く。

2. ノートブックのアクセサレーターをTPUに変更

ランタイム → ランタイムのタイプを変更 → ハードウェアアクセサレータのプルダウンからTPUを選択 → 下部の保存を押す。

3. Google Driveをマウント(使う場合のみ)

from google.colab import drive
drive.mount('/content/drive')

以上をセルに入力すると、Google Drive File Streamを承認してねとurl付きのメッセージが表示されるので、urlを踏み、問題なければ表示されたページで承認する。 すると文字列が出力されるので、コピー。セルに戻ってテキストボックスに貼り付けてエンターを押す。

成功すると、以降ノートブックから"/content/drive/My Drive/"でGoogle Driveにアクセスできるようになる。

4. 外部プログラムのコピー

!pip install sentencepiece
!git clone --recursive https://github.com/iki-taichi/tf-keras-transformer.git

sentencepieceはデフォルトのトークナイザーのために使用。
tf-keras-transformerのgit clone時はtokenizerサブモジュールも同時にcloneするために、--recursiveオプションを付けるので注意。

5. リポジトリのディレクトリに移動

import os
os.chdir('tf-keras-transformer')

6. コーパスの用意

!python src/get_kyoto_corpus.py

「(Wikipedia日英京都関連文書対訳コーパス)https://alaginrc.nict.go.jp/WikiCorpus/index.html」をダウンロード後、 tf-keras-transformer/data内に検証用のkyoto_en_ja_valid.csvと学習用のkyoto_en_ja.csvを2:8の割合で出力する。 それぞれ、1行に1対の入力文、出力文をカンマ区切りで並べた標準的なcsvファイル。

7. 学習用の設定と実行

from src.fitting import FitEnvironment
env = FitEnvironment(
        use_tpu=True,
        batch_size=8,
        input_len=(1024, 1024),
        num_epoch=5,
        output_dir='/content/drive/My Drive/transformer_model',
        data_path=['data/kyoto_en_ja.csv'],
        valid_path=['data/kyoto_en_ja_valid.csv'],
        resume_model_path=None,
        resume_initial_epoch=None,
    )
env.run()

設定の意味:

  • use_tpu: TPUを使い場合はTrue。
  • batch_size: 1系列に複数のサンプルを詰め込むのであまり意味はない。はcolabのTPUはデバイス数に合わせて8。
  • input_len: 入出力の最大トークン数を指定(TPUの場合は固定値が必要)。
  • num_epoch: コーパスを周回する回数。ただし、毎回ランダムにサンプルを詰め込むので1エポックの総数は変動するので厳密にすべてのサンプルを同じ回数学習するわけではない。目安。
  • output_path: モデルの保存先。今回はマウントしたGoogle Driveを指定。
  • data_path: 学習データへのパスのリスト。
  • valid_path: 検証データへのパスのリスト。エポックの終わりに毎回ロスとtoken-wiseの正解率を計算(Noneも可)。
  • resume_model_path: 学習済みのモデルを追加で学習する場合に指定。
  • resume_initial_epoch: 追加学習する場合の始まりのエポック番号を指定(そこからステップ数を算出して、学習率のスケジュールを決める)。

今回はデフォルトを使用しているが、モデルのパラメータを指定する場合はmodel_configを指定。
↓ デフォルトのmodel_config

from src.transformer import Config
model_config = Config(
        src_tokenizer='sp_uncase_en_ja_40000',
        tar_tokenizer='sp_uncase_en_ja_40000',
        use_same_embed=True,
        block_num=(6, 6),
        embed_dim=768,
        hidden_dim=3072,
        head_num=12,
        attention_activation='relu',
        feed_forward_activation='gelu',
        dropout_rate=0.10,
    )

8. 待つ

ログの前半

主にモデルのサマリーが出力される。
今回は、エンコーダー・デコーダーのembedding、デコーダーのpresoftmaxの3つを重み共有し、エンコーダー・デコーダーそれぞれにresidual blookを6個積んで、合計 129,997,888パラメータのモデル (すべて学習可能パラメータ)を学習。

次の図はTensorboardのグラフ出力。

f:id:lang-int:20190429155725p:plain
whole structure

エンコーダーはSelf-Attention/Feedforwardで1単位だがデコーダーはSelf-Attention/Query-Key-Attention/Feedforwardで1単位なので長い。

それぞれが重複して2本出てくる理由は不明。

ログの後半

モデルのコンパイルが終わると学習が始まる。

WARNING:tensorflow:Method (on_train_batch_begin) is slow compared to the batch update (0.988424). Check your callbacks.
step_run=100 step_total=100 step_per_epoch=2909 loss=9.4850 masked_sparse_categorical_accuracy=0.0247 lr=7.0605e-06
step_run=200 step_total=200 step_per_epoch=2909 loss=7.9419 masked_sparse_categorical_accuracy=0.0880 lr=2.1324e-05
step_run=300 step_total=300 step_per_epoch=2909 loss=7.0836 masked_sparse_categorical_accuracy=0.1150 lr=3.5588e-05
...

ログ用のcallbackはcallbacks.TensorBoardとバッチごとに指標を監視するコールバック(NBatchLogger)を入れて、100エポックごとに更新。
はじめだけ、callbackが遅いという警告が出るが、問題はなさそう。

検証データを入れている場合は途中でモデルの再コンパイルが始まり、その後検証用の計算が起きる。 再コンパイルは1回目だけ起きる。

INFO:tensorflow:Finished compiling. Time elapsed: 67.5693416595459 secs
296/296 [==============================] - 147s 498ms/step - loss: 3.8432 - masked_sparse_categorical_accuracy: 0.3716
step_run=3000 step_total=3000 step_per_epoch=2909 loss=3.7968 masked_sparse_categorical_accuracy=0.3626 lr=4.2071e-04
step_run=3100 step_total=3100 step_per_epoch=2909 loss=3.7822 masked_sparse_categorical_accuracy=0.3728 lr=4.3497e-04
...

TensorBoardのグラフ

f:id:lang-int:20190429160042p:plain
transformer_learning_curve

検証データのステップ表記がずれていたので学習データに合わせて伸縮。 もう少し学習・検証ともにロスが下がりそう(5エボックでは短かったと思われる)。

学習に要した時間 合計 6483 s
内訳

学習後に出力されるファイル

  • .hdf5のモデルファイル(5エポックごと + 全部完了後)
  • kerasのhistory callbackの出力(json)
  • tensorboard用のログ(events.out.tfevents.*)

90分を超える場合はリセット対策

Google Colaboratory の 90分ルールを回避したい。 - @Gimina_Graph Qiita

9. 学習済みモデルで推論する

tf.keras.backend.clear_session()
from src.transformer import TransformerWrapper
trans = TransformerWrapper('/content/drive/My Drive/transformer_model/model.05.hdf5')
trans('Kyoto is a Japanese city.')
# '京都(きょうと)は、日本の市。'

変な挙動はしてなさそう。

感想

  • ネットにつながるだけで誰でもTPUで計算できる。

  • 8並列が強い。メモリー容量上GPU1台と乗せられるモデルの大きさはあまり変わらないが、8並列なので(GPUと等速だとしても)最大8倍高速に学習できる。

  • ただし、モデルの転送やコンパイルに時間がかかるのでデータ量が少量の場合は逆に時間がかかる。

以降、細かい部分。

ベースとして利用したtransformerの実装

CyberZHGさんがMITライセンスで公開しているkeras-transformerをベースとして利用。

フレームワークはkeras。

pypiで公開されているが、今回はtensorflow用に書き換える必要があったためpip installではなく、依存するファイルを洗い出してコピーしてきて改変。

主な変更が必要だったこと

tf.kerasで書かれたモデルはtf.contrib.tpu.keras_to_tpu_modelでTPUで動くモデルに変換できる。 これを利用してモデルを変換するため、また、効率的にモデルに入力が流せるようにするためにいくつかの変更を施した。

  • kerasをtf.kerasに変更
  • バッチサイズ以外の次元が確実に数値で入るようにする
  • Embedding層のweight出力前に適当な演算を挟む
  • マスク付きのロスを自前で実装する
  • 固定シーケンス長でも同時に複数サンプルを学習できるようにPos-IdsとGroup-Idsを導入

kerasをtf.kerasに変更

これは、書いてある通り。 keras-transformerの各ファイルの上部に"import tensorflow as tf"を追加し、機械的にkerasをtf.kerasに置換した。
kerasのモデルのままでkeras_to_tpu_modelすると例外がおきる。

バッチサイズ以外の次元が確実に数値で入るようにする

tf.kerasに変更したら、そのまま動くことを期待したが、そうもいかなかった。
0次元(batchsize)以外のテンソルの次元がNoneではダメと怒られるので、まず入力シーケンス長を固定値にする。 また、入力シーケンス長を固定値にしても、途中でreshapeを挟むと次元が決まっているのにプレースホルダーになってしまう箇所がある。そうならないよう数値を指定するように変更(multi_head_attention.pyのここなど)。

Embedding層のweight出力前に適当な演算を挟む

今回のtransformerでは、入力と出力の重みを共有するためにEmbedding層が(普通の埋め込みベクトル、そのもととなるレイヤーの重み全体)の2つを出力するように実装されている。
そして、もととなるレイヤーの重み(variable)をそのまま出力してしまうと、「ReplicatedVariableはkerasのレイヤー出力ではない」というエラーが出てしまう。 色々試した結果、ダミーの演算(tf.cast(weights, tf.float32)など)をレイヤーを出す前に挟んでやることでこのエラーが回避できることが分かった(ここ)。
学習は進んでいるので大丈夫だと思うが、ひょっとしたら、何か副作用があるかもしれない。

マスク付きのロスを自前で実装する

上の3つでとりあえず動作するようになるが、Embeddingのzero_maskをTrueにしてpaddingされている部分のマスクを指定しても、現状のTPUモデルではなぜかlossの計算にマスクが反映されない。 そこで、マスクを有効にするため、標準のsample_weights, maskは使わずに、自前のマスク付きのロスを用意した(masked_softmax_cross_entropy_with_logit)。
なお、このロスは入力側からくるマスクではなく、chainerのように教師データ側の-1ラベルをマスク対象とするように作成した。

固定シーケンス長でも同時に複数サンプルを学習できるようにPos-IdsとGroup-Idsを導入

上の4つで学習までできるようになるが、シーケンス列を固定長にして長い固定長を設定した場合にパディングばっかりになってしまうため、効率が悪い。 そこで、入力にPos-IdsとGroup-Idsを付け足し、1つの系列に複数のサンプルを詰め込むように改造した。
Token-Ids: トークンの識別番号(これまでの入力)
Group-Ids: ひとつの系列に含まれる同一サンプルのトークンに同じ番号を割り当てる。Self-Attention, Query-Key-Attentionをするときに同じ番号のトークンを条件としてマスクをあてることで、他のサンプルの情報を見ないようにする。
Pos-Ids: トークンがサンプル内の何番目かを入力する。

デモとしてダウンロードするコーパスの情報

https://alaginrc.nict.go.jp/WikiCorpus/index.htmlで配布されているWikipedia日英京都関連文書対訳コーパスを使用する(使用時はライセンスを確認のこと)。

国立研究開発法人情報通信研究機構が作成したコーパス。京都に関係あるwikipedia日本語版の記事を人手で翻訳したコーパス(45万文対くらいあったと思います)。

日本語で無料で手に入るデータセットとしては質・量ともにレベルが高いと思われる(ただし、用語と言葉遣いに多少の偏りがある)。

テキストのトークナイズ

SentecePiece。
英語wikipedia, 日本語wikipediaから持ってきた600,000記事(2言語の記事が大体バランスするようにサンプル)の本文で学習したSentecePieceをトークナイザーとして使用。合計で40,000トークン。
トークナイザーは別リポジトリとして作成し、submoduleとして参照。

このトークナイザーでは、スペースを強制的に1トークンにする設定にしたが、英語をトークナイズするときスペースが系列長を無駄に占めてしまうのが若干課題。
(英語の入力長が日本語に比べてかなり長くなってしまう傾向があるため、出力側に無駄が多い)

google/bertリポジトリのFullTokenizerと同一のインターフェースのため、必要であれば取り換え可能。

全体の参考