ELECTRAをTPUで事前学習 [GCP]

GCP経由でTPUを利用することによって，ELECTRAの事前学習を実行する方法を紹介．

必要なサービス

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

手順

Cloud Shell ウィンドウを開く

Cloud Shell (GCPのホーム画面右上のボタンを押すとでてくる)を開き，コマンドラインを出現させる．

画面下半分にCloud Shellが出現

プロジェクトIDの変数を作成

export PROJECT_ID=project-id

GCPではプロジェクト単位で課金されるため，さきにプロジェクトidを入力しておく．

Cloud TPUを作成するプロジェクトを指定

gcloud config set project $PROJECT_ID

先ほどの変数を使って，gcloudコマンドラインツールを構成

このコマンドを初めて実行すると，以下のような画面がでるため，「承認」ボタンを押す．

Cloud Storageバケットを作成

gsutil mb -p ${PROJECT_ID} -c standard -l us-central1 -b on gs://electra

gs://electraは筆者がバケットに割り当てた任意の名前で，他の名前でも構わない．

ここで作成したバケットに入力データを格納する．入力データは大量のテキストをtfrecordに変換したもので，ELECTRAのリポジトリにあるbuild_pretraining_data.pyを利用し作成する．作成の方法は，リポジトリのREADMEを参考にされたし．

また，モデルの学習結果もこのバケットに格納されることを想定している．

インスタンスとTPUを起動

ctpu up --project=${PROJECT_ID} --zone=us-central1-b --tpu-size=v3-8 --tf-version=1.15.3 --name=electra-tpu

ELECTRAの仕様に合わせるため，tf-versionは1.15.3に設定．

ctpu will use the following configuration:

  Name:                 electra-tpu
  Zone:                 us-central1-b
  GCP Project:          xxxx
  TensorFlow Version:   1.15.3
  VM:
      Machine Type:     n1-standard-2
      Disk Size:        250 GB
      Preemptible:      false
  Cloud TPU:
      Size:             v3-8
      Preemptible:      false
      Reserved:         false

OK to create your Cloud TPU resources with the above configuration? [Yn]:

実行後，確認のメッセージがでるため，OKならYを入力．

ctpu upコマンドの実行後，shellプロンプトがusername@projectからusername@tpunameに変更されたことを確認．変更されていれば，Compute Engine VMにログインしていることになる．

ELECTRAのソースコードをgit cloneで入手

Compute Engine VM上で，以下のコマンドを実行．

git clone https://github.com/google-research/electra.git
cd electra

ELECTRAを実行

各種パラメータを設定し，run_pretraining.pyで事前学習開始．GCPサーバーでスクリプトを動かす際はnohup &で挟んでssh接続が切れても動くようにしておくこと．この実行コマンドの場合，同階層にelectra.logが出力され，ここにログが出力される．

nohup python3 run_pretraining.py --data-dir gs://electra/data --model-name my_electra_base --hparams '{"use_tpu":true, "num_tpu_cores":8, "tpu_name":"electra-tpu", "vocab_size":64000, "do_lower_case":false, "num_train_steps":500000, "train_batch_size":512, "generator_hidden_size":1.0, "model_size":"base"}' > electra.log &

SSH接続が切れた後にもう一度TPUが稼働しているインスタンスに接続するときは，以下のコマンドを入力．

gcloud compute ssh electra-tpu --zone=us-central1-b

TPUの削除

学習後は速やかにインスタンスとTPUを停止，あるいは削除すること．稼働中は常に課金されているため，放置しているととんでもない金額が請求されるからだ．

公式の手順によるとコマンドラインから削除できるようだが，筆者はGUI上で，VMインスタンスと，TPUを停止，あるいは削除することをおすすめする．GUIだと間違えて違うインスタンスを削除することを防げたり，TPUのモニタリングを確認してから削除することができるためである．

ELECTRA,Google Cloud Platform,Tensorflow

Posted by vastee