GCP AI Platformで画像認識してみる（学習編）

Google Cloud PlatformのML EngineがAI Platformという名前に変わってリリースされました。

ちょっと僕も触ってみようか、ということでCIFAR-10の画像分類モデル (w/ TensorFlow 2.0a, Python 3.5) をGCPのAI Platformで学習してみました。

クラウド初心者なので回りくどい説明だったり誤記もあるかもしれません。個人用メモの転載ということで容赦してください。

f:id:xterm256color:20190505031523j:plain

目次
ローカル環境の準備
- cloud-sdk/gsutil
GCPプロジェクトの作成
Google Cloud Storage (GCS) バケットの作成
ソースコード（パッケージ）の用意
- 注意点1 GCSとの入出力
- 注意点2 勝手に追加される引数
ジョブの投入
まとめ
補足メモ
- ランタイムバージョンとカスタムイメージ利用について

ローカル環境の準備

cloud-sdk/gsutil

クイックスタートに従ってローカルマシンにcloud-sdkおよびgsutilをインストールしておく。

GCPのCloud Shell機能を使うのであればスキップしても良い。

GCPプロジェクトの作成

適当な名称で作る。

Google Cloud Storage (GCS) バケットの作成

適当な名称でバケットを作り、CIFAR-10(cifar-10-binary.tar.gz)を解凍したファイルをアップロードする。

バケットはマルチリージョンで作るとダメなので、単一リージョンのバケットとして作ることに注意する。

ここでは gs://my-storage/cifar-10/cifar-10-baches-bin の中にフラットにファイルを置いたものとする。

gs://my-storage/cifar-10/cifar-10-batches-bin/batches.meta.txt
gs://my-storage/cifar-10/cifar-10-batches-bin/data_batch_1.bin
gs://my-storage/cifar-10/cifar-10-batches-bin/data_batch_2.bin
gs://my-storage/cifar-10/cifar-10-batches-bin/data_batch_3.bin
gs://my-storage/cifar-10/cifar-10-batches-bin/data_batch_4.bin
gs://my-storage/cifar-10/cifar-10-batches-bin/data_batch_5.bin
gs://my-storage/cifar-10/cifar-10-batches-bin/readme.html
gs://my-storage/cifar-10/cifar-10-batches-bin/test_batch.bin

ソースコード（パッケージ）の用意

AI Platformでは、あらかじめPythonパッケージとしてソースコードを用意しておく。（PyPIにあげたりする必要はない）

以下のようなプロジェクト構造が推奨されてるようだが、普通のPythonパッケージのルールに従っていれば特に問題はない。

MyProject/
    trainer/
        __init__.py
        task.py     # トレーニングジョブを管理するアプリケーションロジック
        model.py    # モデルのロジックである TensorFlow グラフコード
        util.py     # （存在する場合）には、トレーニング アプリケーションを実行するためのコード
    setup.py

（トレーニングアプリケーションのパッケージング | TensorFlow 用 Cloud ML Engine | Google Cloudより抜粋）

ここでは以下のソースコードを使用した。

https://github.com/chmod644/cifar-10-ai-platform

こんなプロジェクト構造になっている。

cifar-10-ai-platform/
    cifar_10/
        __init__.py
        train.py    # トレーニングジョブを管理するアプリケーションロジック
        model.py    # モデルのロジックである TensorFlow グラフコード
        input.py    # 入力ロジック
    setup.py

エントリポイントであるtrain.pyの内容としては、cifar-10-batchesで10クラス分類するモデルを学習して、学習済みモデルをSavedModel形式で保存するようなコードになっている。

なお、当然のことながらローカルで学習する場合といくつか実行環境が違うので意識しないといけない点がある。特に以下の2点は実装に影響するところなのでよく注意する必要がある。

注意点1 GCSとの入出力

AI Platformを使う場合は基本的にはGCS(もしくは他のクラウドストレージなど)との間でファイル入出力するように書くはずだが、これを意識してないと確実にコケる。

TensorFlowを使っているならtf.io.gfileというAPIを使うのがおすすめ。ファイルオープンやglobなんかを透過的に扱えるようになっている。（TensorFlow以外でも使えるように別パッケージにしたらいいのにと思う）

注意点2 勝手に追加される引数

ジョブを投入すると、AI Platformがシステム的に--job-dirという引数を追加する。この引数を受けれるようにスクリプトを作っておくか、未定義の引数も受けれるように作っておく必要がある。

未定義の引数を受け取る方法は以下。 - argparseなら parser.parse_known_args() で引数解析する - absl-pyなら実行時に --undefok=job-dir を末尾に追加する

ジョブの投入

ジョブの投入は、gcloudが使える環境下で実行する。

マシンタイプの設定ファイル

GPU対応マシンを使ったりするにはconfig.yamlにマシンタイプを明記して、gcloudコマンドに渡すことにする。今回はGPU1枚だけを使う設定としてBASIC_GPUを選択する。

trainingInput:
  scaleTier: BASIC_GPU

この設定ファイルで分散学習用のマシンタイプ指定とかもできる。詳しくは → REST Resource: projects.jobs | AI Platform | Google Cloud

補遺：gcloud betaであればコマンドラインフラグに指定できる

ジョブの投入コマンド

あとはgcloudコマンドでジョブを実際に投入するだけ。

Runtime Version 1.13でTensorFlow 1.13に対応してるようなのでこれを使ってみることにする（Runtime Versionについては後述）。ちなみに1.12以下を使うとCUDAのバージョン不整合で落ちた。

Runtime Version List | AI Platform for TensorFlow | Google Cloud

# パッケージの下に移動
git clone https://github.com/chmod644/cifar-10-ai-platform
cd cifar-10-ai-platform

# 環境変数の設定
now=$(date +"%Y%m%d_%H%M%S")
JOB_NAME=cifar_10_$now  # ジョブ名はプロジェクト内で一意な名称を付ける必要がある
REGION=us-central1      # バケットのリージョンと一緒にしておく
BUCKET_NAME=my-storage
MODEL_DIR=output
INPUT_DIR=gs://$BUCKET_NAME/cifar-10/cifar-10-batches-bin/
OUTPUT_PATH=gs://$BUCKET_NAME/cifar-10/$JOB_NAME

# ジョブの投入
gcloud ai-platform jobs submit training $JOB_NAME \
    --job-dir $OUTPUT_PATH \
    --runtime-version 1.13 \    # tf2.0のGPU版を使うため1.13を指定
    --python-version 3.5 \      # Python 2.7を使う場合は不要
    --module-name cifar_10.train \  # エントリポイント
    --package-path cifar_10 \ 
    --region $REGION \
    --config config.yaml \      # 後述
    -- \                        # これ以降はcifar_10.trainの入力
    --input $INPUT_DIR \
    --output $OUTPUT_PATH \
    --epochs 1 \
    --undefok=job-dir

実行に成功すると以下のようなメッセージが表示される

Job [cifar_10_20190505_005847] submitted successfully.
Your job is still active. You may view the status of your job with the command

  $ gcloud ai-platform jobs describe cifar_10_20190505_005847

or continue streaming the logs with the command

  $ gcloud ai-platform jobs stream-logs cifar_10_20190505_005847
jobId: cifar_10_20190505_005847
state: QUEUED

動作確認

AI Platformのジョブ画面からログや実行状況が確認できる。（Kerasのmodel.fitは出力が長くなるので、オフにしたほうがよかったかも）

また、GCS内にSavedModelが出力されていることが確認できる。

gs://$BUCKET_NAME/cifar-10/$JOB_NAME/
├── cifar-10-model
│   └── 1556986596
│       ├── saved_model.pb
│       └── variables
│           ├── variables.data-00000-of-00001
│           └── variables.index
├── log
│   ├── flags.json
│   └── model.txt
└── packages
    └── 45d583aa04be9064f4b9f787b733264ac375029243dca3eabf050ccce1536c49
        └── cifar_10-0.1.tar.gz

logやcifar-10-modelはcifar10.trainが出力したもの。packagesだけはシステムが追加したもので、実行時点のパッケージ（cifar-10-ai-platform）が格納されている。

まとめ

とりあえずGPUを使った学習をまわしてSavedModelを出力するとこまでした。

他にも下記のあたりを試したいですが、眠いのでまた後日調査します。

学習済みモデルのデプロイ方法とクライアントからのアクセス
Cloud Dataflowを利用した前処理の自動化
AI Platform内のNotebook機能の使い勝手
Cloud Shellの永続的カスタマイズ

補足メモ

ランタイムバージョンとカスタムイメージ利用について

AI Platformの学習ジョブや推論ジョブはイメージを使用して学習・推論を行うVMを構成する。そのイメージにあたるものが「ランタイムバージョン」と呼ばれている。

どのランタイムバージョンにどのパッケージが入ってるのか、またそのパッケージバージョンは何なのかは以下のページで調べることができる。