端の知識の備忘録

技術メモになりきれない、なにものか達の供養先

KaggleでGoogle ColabやTPUを使うためのTips集

概要

意外とKaggleでColabやTPUを使う情報が少ないので、Kaggle NotebookをTPU対応させたりそれをColab対応させたりするときのTipsをご紹介します。基本はTensorflowに関する記述です。

strategyを定義してmodel.fitのときにstrategy.scope()のブロックで実行するみたいな、基本的な使い方は公式チュートリアル参照のこと(ぶっちゃけ自分もまだこの辺よくわかってなくておまじない的に使ってますが)。 https://cloud.google.com/tpu/docs/tutorials?hl=ja

エラーコードなども随時補完していきたいと思います。

KaggleデータセットをColabから使う方法

私も使うまで知らなかったのですが、TPUを使いたいときには、GCSにTFrecordのような分散させた形で配置したデータセットを用意しないと使えません！。正確にはTFrecordでなくてもGCSに画像を置くだけでもTPUを使ったデータの読み込み＆Trainは可能ですが、まともに速度が出ません。

Kaggle datasetは実態としてGCSへ格納されているらしく、GCSのパスを取得することができればこれをColabから利用可能です。

これから先ではPublicなKaggle datasetをColabで使う手順を紹介しますが、Privateなデータセットも認証とか頑張れば行けるらしい。

  • kaggleノートブックでGCSパスを取得

PublicなKaggleデータセットをColabで使うには、KaggleデータセットのGCSパスをKaggleノートブックで取得して、それをColab側で使うことでアクセス可能です。

参考： https://www.kaggle.com/code/tchaye59/efficientnet-tensorflow-baseline-tpu/notebook

使いたいKaggleデータセットをInputにしたノートブックで、下記コードを実行。

from kaggle_datasets import KaggleDatasets
GCS_PATH = KaggleDatasets().get_gcs_path()
print(GCS_PATH)

# 以下出力
gs://kds-c6262c3899ad974e1123f2b7bad13b081bfad9d034edac5d4cfad2bf

これで取得したGCS_PATHを何処かにメモっておき、これをColab側で使います。

  • Colab側での利用方法

参考: https://zenn.dev/karunru/articles/a8f394d7a859b7

Kaggle APIがどうのこうのみたいなしゃらくさい話はさておいて、とりあえずKaggle datasetをColabで利用するための(おそらく)最短手順は下記のとおりです。

GCS_BUCKETNAME = "さっきKaggleノートブックで確認したGCS_PATHから'gs://'を除いた部分"

# パブリックデータセットであっても、authenticationを行う必要があるらしい。
from google.colab import auth
auth.authenticate_user()

# TFDSのtfds.core.GeneratorBasedBuilderに従って作られたデータセットであれば、data_dirとしてGCS_PATHを指定すれば読み込める。読み込めないときは、PATHの深さがずれていないか確認すると良い。
# 例として示しているFGVCDataset は、次のノートブックを参考に作ったTFrecord形式のデータセット。 https://www.kaggle.com/code/tchaye59/efficientnet-tensorflow-baseline-tpu/notebook

data_dir= "gs://" + GCS_BUCKETNAME
builder = FGVCDataset(data_dir=data_dir)

# もしマウントする必要がなければ下記不要。
!echo "deb http://packages.cloud.google.com/apt gcsfuse-bionic main" > /etc/apt/sources.list.d/gcsfuse.list
!curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
!apt update
!apt install gcsfuse
!mkdir -p /kaggle/input
!gcsfuse --implicit-dirs --limit-bytes-per-sec -1 --limit-ops-per-sec -1 " "/kaggle/input"

# これで'/kaggle/input'下にKaggle datasetがマウントされます。
!ls /kaggle/input

ちなみに、一度だけ遭遇したのですが、Kaggle datasetのGCS_PATHが勝手に変わった事があり、それで学習がストップしたことがありました。そういうこともあると思って、もう一度GCS_PATHを取得し直しましょう。

TPU使用時の豆知識

TPUを使っていると、同じコードでも謎にエラーを出すことがあるので、そういうときは慌てずもう一度ランさせましょう。何故か普通に動きます。

例えばこんなエラー。

AlreadyExistsError                        Traceback (most recent call last)
<ipython-input-12-effe1bc08a91> in <module>()
      1 # https://stackoverflow.com/questions/59289014/how-to-check-if-tpu-device-type-is-v2-or-v3
      2 tpu_worker = os.environ['COLAB_TPU_ADDR'].replace('8470', '8466')
----> 3 print(tf.profiler.experimental.client.monitor(tpu_worker,1))
      4 tpu_info = tf.profiler.experimental.client.monitor(tpu_worker,1)
      5 tpu_type = tpu_info.split('\n')[1].split(':')[1].split(" ")[2]

/usr/local/lib/python3.7/dist-packages/tensorflow/python/profiler/profiler_client.py in monitor(service_addr, duration_ms, level)
    162   """
    163   return _pywrap_profiler.monitor(
--> 164       _strip_prefix(service_addr, _GRPC_PREFIX), duration_ms, level, True)
    165 
    166 

AlreadyExistsError: Another profiling session is active

model saveは重みだけをh5で保存する必要あり

TPUで学習させたモデルは、重みだけ保存する必要があるようです。ですので、

model.save("MODEL_NAME.h5")

のように、拡張子をh5にして保存を行います。

TPUのバージョン確認方法

これが意外とどこにも書いてなかった。

Google ColabだとTPUのバージョンを指定できないので、V2とV3でBatchサイズを変えるみたいな処理をしたいことがあると思います(Colab Proだと実際にはV2しか割り当てられなさそうですが……)。そういうときには次のようなやり方でTPUの情報を取得します。

TPU接続後に下記コードを実行。

参考： https://stackoverflow.com/questions/59289014/how-to-check-if-tpu-device-type-is-v2-or-v3

tpu_worker = os.environ['COLAB_TPU_ADDR'].replace('8470', '8466') 
print(tf.profiler.experimental.client.monitor(tpu_worker,1))
tpu_info = tf.profiler.experimental.client.monitor(tpu_worker,1)

# 以下出力
  Timestamp: 13:15:07
  TPU type: TPU v2
  Utilization of TPU Matrix Units (higher is better): 0.000%

ちなみにここで取得したtpu_infoを次のように処理(だいぶ雑ですが)すると、v2またはv3の文字列を取得できるので、これで処理を分けることができます。

tpu_type = tpu_info.split('\n')[1].split(':')[1].split(" ")[2]

# 例えばこんな感じ
if tpu_type == "v2":
  BATCH_SIZE = 512
else:
  BATCH_SIZE = 1024

TPUのパフォーマンス系Tips

参考： https://cloud.google.com/tpu/docs/troubleshooting/trouble-tf?hl=ja#overview

IOがネックになることが多いので、それを回避するあれこれ。

  • BATCHSIZE

バッチサイズは64の倍数にすべし。TPUでは内部的には8個のコアに分散して処理を行っているため、バッチサイズ64ならそれぞれのコアがバッチサイズ64/8＝8で学習を行う。

それぞれのコアはTPU v2なら8GB、TPU v3なら16GBのメモリを持っているので、結構大きなサイズで学習させられます。Googleのガイド的には1024から試し始めて、OOMになったら下げるというやり方でBatchサイズを調整するようにアドバイスされています。

TPU利用時にはモデルのコンパイル時にsteps_per_executionという引数を指定することができ、指定した数だけまとめてトレーニングステップを実行してから、TPUからのコールバック(lossのロギングやチェックポイントの保存)がまとめて行われるというもの。

つまり、学習中の

64/346 [====>.........................] - ETA: 2:29 - loss: 0.2038 - accuracy: 0.9377

みたいな表示の更新が、steps_per_executionぶんだけまとめて行われ反映が遅くなり、そのタイミングでチェックポイントの保存も行われることになります。

プログレスバーが常時更新され伸びていくのを見れなくなるのは残念ですが、それ以外特に大きな欠点はないと思うので(要出典)、とりあえず32とか設定しておくとTPUでの学習が速くなると思います。

model.compile(optimizer=tf.keras.optimizers.Adam(), 
                loss='sparse_categorical_crossentropy',
                metrics=["accuracy"],
                steps_per_execution=steps)