5分でできる：Googleの自然言語処理AI（BERT）をTPU上で転移学習

BERTは、種類の異なる様々な自然言語処理タスクで最高成績を叩きだした優れモノで、
しかも転移学習のベースにすれば少ない学習のみで個々の問題を従来より正確に解ける可能性がある「革命的」な技術。

• Googleのオリジナルコードを元にしていますが、処理の順序をわかりやすく解説する為に若干変更しています。
• 無料でCloud TPUが使用できるGoogle Colab環境（notebook形式）を使用します。
• 本記事で使用したnotebookはココ。
• GCS（Googleクラウドストレージ）のアカウントがあれば5分で実行できます。
• 転移学習に使用しているMRPCデータセットで何？というのは別の方の記事を参照

参考：TPUは、GPUより4倍ほど高速でした。（Google Colab環境で比較）

計算機	転移学習の時間
TPU	3分
GPU ※1	12分
CPU	374分

※1: Colab環境のGPUは、Tesra K80（GPUコアを2つ搭載）の内のGPUコア1個分です

手順解説

１、まず、google.colab.authモジュール中のauthenticate_user()を使用して、Googleサービスにログインします。

この処理は、後述のrun_classifierスクリプトを実行する前に実行が必要です。リンクをクリックして開くGoogleのページでログインを行い、表示される文字列を下記の入力欄にコピーしてEnterします。

from google.colab import auth
auth.authenticate_user()

!ls .

成功すると、ローカルのホームディレクトリにadc.jsonというファイルが生成されます。

adc.json

２、次に、githubよりBERTのレポジトリをクローンし、パスを通しておきます。

この処理は、Googleが提供しているBERTのライブラリを後ほどimportする為に必要です。

#testコマンドを実行するとbert_repoというgitレポジトリがcloseされる
!test -d bert_repo || git clone https://github.com/google-research/bert bert_repo
if not 'bert_repo' in sys.path:
  sys.path += ['bert_repo']
  
!ls -a bert_repo

下記のファイルがダウンロードできました。

.			    LICENSE		  requirements.txt
..			    modeling.py		  run_classifier.py
CONTRIBUTING.md		    modeling_test.py	  run_pretraining.py
create_pretraining_data.py  multilingual.md	  run_squad.py
extract_features.py	    optimization.py	  sample_text.txt
.git			    optimization_test.py  tokenization.py
.gitignore		    __pycache__		  tokenization_test.py
__init__.py		    README.md

３、転移学習の為のデータを準備します。

• ここでは、MRPCデータセットの転移学習を行うために、「GLUEからMRPCのデータをダウンロードする為のスクリプト」をgithubからクローンし、ダウンロードを実行します。
• ダウンロードしたファイルは、Google提供の実験用スクリプト（run_classifierモジュール）を使用して読み込みます。

# Specify training task
TASK = 'MRPC'

# Download glue data.
! test -d download_glue_repo || git clone https://gist.github.com/60c2bdb54d156a41194446737ce03e2e.git download_glue_repo
!python download_glue_repo/download_glue_data.py --data_dir='glue_data' --tasks=$TASK
TASK_DATA_DIR = 'glue_data/' + TASK
print('***** Task data directory: {} *****'.format(TASK_DATA_DIR))
!ls $TASK_DATA_DIR

print('***** load data *****')
import run_classifier

processor = run_classifier.MrpcProcessor()
label_list = processor.get_labels()
train_examples = processor.get_train_examples(TASK_DATA_DIR)
eval_examples = processor.get_dev_examples(TASK_DATA_DIR)
print("label_list",label_list)
print("train_examples",len(train_examples))
print("eval_examples",len(eval_examples))

• ダウンロードが成功していれば、訓練用データの数が3668、開発評価用データの数が408と出力されるはずです。

Processing MRPC...
	Completed!
***** Task data directory: glue_data/MRPC *****
dev_ids.tsv  msr_paraphrase_test.txt   test.tsv
dev.tsv      msr_paraphrase_train.txt  train.tsv
***** load data *****
label_list ['0', '1']
train_examples 3668
eval_examples 408

４、転移学習のベースとする事前学習済みBERTモデルを準備します。

• ここでは、入力文を小文字化する版でモデルサイズが小さい（12層／768隠れ層）方のモデルをGCS（Googleクラウドストレージ）からダウンロードします
• 転移学習に使用するバッチサイズやエポック数等のパラメータを指定し、モデルを構築（する為の関数を作成）します。モデルを構築する為の関数は、Google提供の実験用スクリプト（run_classifierモジュール）を使用することで簡単に作成できます。use_tpu=TrueとすることでTPU上で実行が可能になります。ColabのランタイムのタイプがTPUでない（CPUやGPU）場合はエラーになりますので、ランタイムタイプをTPUに変更して、最初から再実行してください。

# Specify BERT pretrained model
BERT_MODEL = 'uncased_L-12_H-768_A-12' #@param {type:"string"}
BERT_PRETRAINED_DIR = 'gs://cloud-tpu-checkpoints/bert/' + BERT_MODEL
print('***** BERT pretrained directory: {} *****'.format(BERT_PRETRAINED_DIR))
!gsutil ls $BERT_PRETRAINED_DIR

import modeling

# Train configs
TRAIN_BATCH_SIZE = 32 #@param {type:"integer"}
EVAL_BATCH_SIZE = 8 #@param {type:"integer"}
TRAIN_EPOCHS = 3.0 #@param {type:"number"}
WARMUP_PROPORTION = 0.1 #@param {type:"number"}
LERANING_RATE = 2e-5 #@param {type:"number"}

train_steps = int(len(train_examples) / TRAIN_BATCH_SIZE * TRAIN_EPOCHS)
eval_steps = int(len(eval_examples) / EVAL_BATCH_SIZE)

bert_config=modeling.BertConfig.from_json_file(os.path.join(BERT_PRETRAINED_DIR, 'bert_config.json'))

assert 'COLAB_TPU_ADDR' in os.environ, 'ERROR: TPU runtime should be selected'
  
model_fn = run_classifier.model_fn_builder(
    bert_config=bert_config,
    num_labels=len(label_list),
    init_checkpoint=os.path.join(BERT_PRETRAINED_DIR, 'bert_model.ckpt'),
    learning_rate=LERANING_RATE,
    num_train_steps=train_steps,
    num_warmup_steps=int(train_steps * WARMUP_PROPORTION),
    use_tpu=True,
    use_one_hot_embeddings=True)

今回使用したBERT事前学習済みモデルは下記です。
なおgs://は、GCS（Googleクラウドストレージ）を指すスキーマです。

***** BERT pretrained directory: gs://cloud-tpu-checkpoints/bert/uncased_L-12_H-768_A-12 *****
gs://cloud-tpu-checkpoints/bert/uncased_L-12_H-768_A-12/bert_config.json
gs://cloud-tpu-checkpoints/bert/uncased_L-12_H-768_A-12/bert_model.ckpt.data-00000-of-00001
gs://cloud-tpu-checkpoints/bert/uncased_L-12_H-768_A-12/bert_model.ckpt.index
gs://cloud-tpu-checkpoints/bert/uncased_L-12_H-768_A-12/bert_model.ckpt.meta
gs://cloud-tpu-checkpoints/bert/uncased_L-12_H-768_A-12/checkpoint
gs://cloud-tpu-checkpoints/bert/uncased_L-12_H-768_A-12/vocab.txt

５、転移学習を実行する設定を行います。

• 転移学習後のモデルファイルを保存するための準備をします。cloud TPUを使用する場合、TPUからの出力はGCS（Googleクラウドストレージ）上である必要があります。予めバケットを作成（もしくは既存のバケットを指定）して、下記のOUTPUT_BUCKET変数をそのバケット名に変更してから、下記コードを実行してください。
• TPU上での実行の為に、TPUConfigを生成しています。意味を理解するには、Google CloudのTPUの説明を読む必要があります。
  • Colabで使用できるTPUはv2なので、per_host_input_for_trainingにはPER_HOST_V2を指定。
  • iterations_per_loopには、 TPUからCPUに制御を戻す頻度をステップ数の数で指定します。とりあえず、チェックポイントを保存するステップ数と同じ値を設定しておけばOK。
  • num_shardsには、モデル並列を用いない場合、TPUのコア数を指定すればよいようです（※最新のドキュメントだと、Deprecated, ignored by TPUEstimatorとなっているので、指定する必要はなくなった模様）
    • レプリカ（Replica）とは： Cloud TPUは、2つのTPUコアをもつチップ4つから構成されています（1Cloud TPU = 8TPUコア）。よって、Cloud TPUを効率的に使うには、8コアをそれぞれ活用するようにプログラムを実行する必要があります。TPUEstimatorは、複製された計算（replicated computation）を実行する為の計算グラフを構築し実行します。レプリカ（replica）は、訓練グラフの複製であり、Estimatorに指定したバッチサイズの1/8のサイズのミニバッチをそれぞれのTPUコア上で実行し訓練します。

# prepare output access
TPU_ADDRESS = 'grpc://' + os.environ['COLAB_TPU_ADDR']
print('TPU address is', TPU_ADDRESS)

with tf.Session(TPU_ADDRESS) as session:
  # Upload credentials to TPU for GCS bucket usage. credentials are set for all future sessions on this TPU.
  with open('/content/adc.json', 'r') as f:
    auth_info = json.load(f)
  tf.contrib.cloud.configure_gcs(session, credentials=auth_info)
  print(tf.get_default_graph().get_operations()) #configure_gcsによってGraphにoperationが追加されるっぽい

OUTPUT_BUCKET = 'dl_dev' #@param {type:"string"}
assert OUTPUT_BUCKET, 'Must specify an existing GCS bucket name'
OUTPUT_DIR = 'gs://{}/bert/models/{}'.format(OUTPUT_BUCKET, TASK)
tf.gfile.MakeDirs(OUTPUT_DIR)
print('***** Model output directory: {} *****'.format(OUTPUT_DIR))  

# Run configs
SAVE_CHECKPOINTS_STEPS = 1000 #@param {type:"integer"}

run_config = tf.contrib.tpu.RunConfig(
    cluster=tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS),
    model_dir=OUTPUT_DIR,
    save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS,
    tpu_config=tf.contrib.tpu.TPUConfig(
        iterations_per_loop=SAVE_CHECKPOINTS_STEPS,
        num_shards=8,
        per_host_input_for_training=tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2))

TPUからGCSにアクセスする為の情報がGraphに追加されています（GcsConfigureCredentialsというoperation）

TPU address is grpc://10.71.146.194:8470
[<tf.Operation 'Placeholder' type=Placeholder>, <tf.Operation 'GcsConfigureCredentials' type=GcsConfigureCredentials>, <tf.Operation 'Placeholder_1' type=Placeholder>, <tf.Operation 'GcsConfigureCredentials_1' type=GcsConfigureCredentials>, <tf.Operation 'Placeholder_2' type=Placeholder>, <tf.Operation 'GcsConfigureCredentials_2' type=GcsConfigureCredentials>]
***** Model output directory: gs://dl_dev/bert/models/MRPC *****

６、最後にTPUEstimatorを作成して準備完了です。

• TPUEstimatorは、TensorFlowにおいてグラフの構築・実行（学習・評価・予測の各アクション）をカプセル化した高レベルAPIであるEstimatorのTPU対応版です。

tf.logging.set_verbosity(tf.logging.WARN)

estimator = tf.contrib.tpu.TPUEstimator(
    use_tpu=True,
    model_fn=model_fn,
    config=run_config,
    train_batch_size=TRAIN_BATCH_SIZE,
    eval_batch_size=EVAL_BATCH_SIZE)

• 下記のWARNINGがでますが、実行に影響ありません。

WARNING:tensorflow:Estimator's model_fn (<function model_fn_builder.<locals>.model_fn at 0x7f7eeccd17b8>) includes params argument, but params are not passed to Estimator.

７、入力データの前処理を行います。

• Google提供のBERTライブラリ（tokenizationモジュール）中のFullTokenizerクラスを用いて入力データのトークン化を行います。
• 結果を用いて、ニューラルネットワークにデータを入力する為の関数を定義します。関数の定義には、Google提供の実験用スクリプト（run_classifierモジュール）を使用しています。
• MAX_SEQ_LENGTHは512以下の値を選択してください。

tf.logging.set_verbosity(tf.logging.INFO)

# pre-process input data
MAX_SEQ_LENGTH = 128 #@param {type:"integer"}

# tokenizing
import tokenization

tokenizer = tokenization.FullTokenizer(
    vocab_file=os.path.join(BERT_PRETRAINED_DIR, 'vocab.txt'),
    do_lower_case=BERT_MODEL.startswith('uncased'))

train_features = run_classifier.convert_examples_to_features(
    train_examples, label_list, MAX_SEQ_LENGTH, tokenizer)
eval_features = run_classifier.convert_examples_to_features(
    eval_examples, label_list, MAX_SEQ_LENGTH, tokenizer)

# define input functions
train_input_fn = run_classifier.input_fn_builder(
    features=train_features,
    seq_length=MAX_SEQ_LENGTH,
    is_training=True,
    drop_remainder=True)
eval_input_fn = run_classifier.input_fn_builder(
    features=eval_features,
    seq_length=MAX_SEQ_LENGTH,
    is_training=False,
    drop_remainder=True)

８、TPUEstimatorを使用して転移学習を実行します。

# Train the model
tf.logging.set_verbosity(tf.logging.WARN)

print('***** Started training at {} *****'.format(datetime.datetime.now()))
print('  Num examples = {}'.format(len(train_examples)))
print('  Batch size = {}'.format(TRAIN_BATCH_SIZE))
print('  train_steps = {}'.format(train_steps))
estimator.train(input_fn=train_input_fn, max_steps=train_steps)
print('***** Finished training at {} *****'.format(datetime.datetime.now()))

デフォルトのパラメータでは、凡そ３分で完了します。

***** Started training at 2019-01-04 11:47:48.796496 *****
  Num examples = 3668
  Batch size = 32
  train_steps = 343
***** Finished training at 2019-01-04 11:50:25.272631 *****

９、TPUEstimatorを使用して転移学習の結果を開発評価用データセットを用いて評価します。

# Eval the model
tf.logging.set_verbosity(tf.logging.WARN)

print('***** Started evaluation at {} *****'.format(datetime.datetime.now()))
print('  Num examples = {}'.format(len(eval_examples)))
print('  Batch size = {}'.format(EVAL_BATCH_SIZE))

result = estimator.evaluate(input_fn=eval_input_fn, steps=eval_steps)
print('***** Finished evaluation at {} *****'.format(datetime.datetime.now()))
print("***** Eval results *****")
print('  {} = {}'.format("eval_accuracy", str(result["eval_accuracy"])))
print('  {} = {}'.format("eval_loss", str(result["eval_loss"])))

学習の度に性能はランダムに異なりますが、凡そ84%から88%の間の正解率になるはず。

***** Started evaluation at 2019-01-04 11:50:52.704990 *****
  Num examples = 408
  Batch size = 8
***** Finished evaluation at 2019-01-04 11:51:30.472141 *****
***** Eval results *****
  eval_accuracy = 0.85784316
  eval_loss = 0.7110055

蛇足）

同じ転移学習をCPUで実行したら、6時間14分かかりました。

tf.logging.set_verbosity(tf.logging.WARN)

CPU_OUTPUT_DIR = './bert/models/{}'.format(TASK)

cpu_model_fn = run_classifier.model_fn_builder(
    bert_config=bert_config,
    num_labels=len(label_list),
    init_checkpoint=os.path.join(BERT_PRETRAINED_DIR, 'bert_model.ckpt'),
    learning_rate=LERANING_RATE,
    num_train_steps=train_steps,
    num_warmup_steps=int(train_steps * WARMUP_PROPORTION),
    use_tpu=False,
    use_one_hot_embeddings=True)

cpu_run_config = tf.contrib.tpu.RunConfig(
    model_dir=CPU_OUTPUT_DIR,
    save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS)

cpu_estimator = tf.contrib.tpu.TPUEstimator(
    use_tpu=False,
    model_fn=cpu_model_fn,
    config=cpu_run_config,
    train_batch_size=TRAIN_BATCH_SIZE,
    eval_batch_size=EVAL_BATCH_SIZE)

# Train the model

print('***** Started training at {} *****'.format(datetime.datetime.now()))
cpu_estimator.train(input_fn=train_input_fn, max_steps=train_steps)
print('***** Finished training at {} *****'.format(datetime.datetime.now()))

# Eval the model
print('***** Started evaluation at {} *****'.format(datetime.datetime.now()))
result = cpu_estimator.evaluate(input_fn=eval_input_fn, steps=eval_steps)
print('***** Finished evaluation at {} *****'.format(datetime.datetime.now()))
print("***** Eval results *****")
print('  {} = {}'.format("eval_accuracy", str(result["eval_accuracy"])))
print('  {} = {}'.format("eval_loss", str(result["eval_loss"])))

WARNING:tensorflow:Estimator's model_fn (<function model_fn_builder.<locals>.model_fn at 0x7f7ed128f268>) includes params argument, but params are not passed to Estimator.
WARNING:tensorflow:eval_on_tpu ignored because use_tpu is False.
***** Started training at 2019-01-04 11:59:27.851376 *****
***** Finished training at 2019-01-04 18:13:13.966963 *****
***** Started evaluation at 2019-01-04 18:13:13.967603 *****
***** Finished evaluation at 2019-01-04 18:19:06.805503 *****
***** Eval results *****
  eval_accuracy = 0.8627451
  eval_loss = 0.4576749