Cloud TPU（TF 2.x）での MNIST の実行

このチュートリアルでは、MNIST モデルの概要、MNIST TensorFlow TPU コードサンプルのダウンロード方法、Cloud TPU でコードを実行する方法について説明します。

免責

このチュートリアルでは、サードパーティのデータセットを使用します。Google は、このデータセットの有効性またはその他の側面について、いかなる表明、保証もいたしません。

モデルの説明

MNIST データセットには、0 から 9 までの手書きの数字の画像と、各画像の数字を識別するラベルが含まれています。

このチュートリアルでは、MNIST データセットに基づいて画像を分類する機械学習モデルをトレーニングします。トレーニング後、このモデルは MNIST データセットの手書き画像から学習した内容に基づいて、受信画像を 10 個のカテゴリ（0〜9）のいずれかに分類します。まだ学習していない画像をモデルに送信した場合、モデルはトレーニングで学習した内容に基づいて画像内の数字を識別します。

MNIST データセットは次の 3 つの部分から構成されます。

  • 60,000 例のトレーニング データ
  • 10,000 例のテストデータ
  • 5,000 例の検証データ

モデルには 7 つの層が混在しています。

  • 2 つの畳み込み層
  • 2 つの最大プーリング増
  • 2 つの高密度（全結合）層
  • 1 つのドロップアウト層

損失は、多クラス交差エントロピーによって計算されます。

このバージョンの MNIST モデルでは、Keras API を使用しています。これは、Cloud TPU で機械学習モデルを作成して実行する場合に推奨される方法です。

Keras を使用すると、低レベルの実装の大部分が隠されることで、モデル開発プロセスが簡素化され、TPU と他のプラットフォーム（GPU や CPU など）の切り替えが簡単になります。

目標

  • データセットとモデルの出力を格納する Cloud Storage バケットを作成します。
  • トレーニング ジョブを実行します。
  • 出力結果を確認します。

費用

このドキュメントでは、Google Cloud の次の課金対象のコンポーネントを使用します。

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。 新しい Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。

始める前に

このセクションでは、Cloud Storage バケットと Compute Engine VM の設定について説明します。

  1. Cloud Shell ウィンドウを開きます。

    Cloud Shell を開く

  2. プロジェクト ID の変数を作成します。

    export PROJECT_ID=project-id

  3. Cloud TPU を作成するプロジェクトを使用するように Google Cloud CLI を構成します。

    gcloud config set project ${PROJECT_ID}

    このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した Google Cloud API の呼び出しを許可します。

  4. Cloud TPU プロジェクトのサービス アカウントを作成します。

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID

    このコマンドでは、Cloud TPU サービス アカウントを次の形式で返します。 

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com

  5. 次のコマンドを使用して Cloud Storage バケットを作成します。

    gsutil mb -p ${PROJECT_ID} -c standard -l us-central1 gs://bucket-name

    この Cloud Storage バケットには、モデルのトレーニングに使用するデータとトレーニング結果が格納されます。このチュートリアルで使用する gcloud コマンドは、前の手順で設定した Cloud TPU サービス アカウントのデフォルトの権限を設定します。権限の詳細な設定が必要な場合は、アクセスレベル権限をご覧ください。

  6. gcloud コマンドを使用して Compute Engine VM と Cloud TPU を起動します。使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。詳細については、システム アーキテクチャをご覧ください。

    TPU VM

    $ gcloud compute tpus tpu-vm create mnist-tutorial \
--zone=us-central1-b \
--accelerator-type=v3-8 \
--version=tpu-vm-tf-2.16.1-pjrt \
--preemptible

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン。
    accelerator-type
    アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。
    version
    Cloud TPU ソフトウェアのバージョン。
    preemptible
    プリエンプティブル TPU は低コストですが、いつでもシャットダウンされる可能性があります。

    TPU ノード

    $ gcloud compute tpus execution-groups create \
--name=mnist-tutorial \
--zone=us-central1-b \
--tf-version=2.12.0 \
--machine-type=n1-standard-1 \
--accelerator-type=v3-8 \
--preemptible

    コマンドフラグの説明

    name
    作成する Cloud TPU の名前。
    zone
    Cloud TPU を作成するゾーン。
    tf-version
    gcloud コマンドが VM にインストールする TensorFlow のバージョン。
    machine-type
    作成する Compute Engine VM のマシンタイプ。
    accelerator-type
    作成する Cloud TPU のタイプ。
    preemptible
    プリエンプティブル TPU は低コストですが、いつでもシャットダウンされる可能性があります。

    gcloud コマンドの詳細については、gcloud リファレンスをご覧ください。

  7. gcloud compute tpus コマンドの実行が終了したら、shell プロンプトが username@projectname から username@vm-name に変更されたことを確認します。変更されていれば、Compute Engine VM にログインしていることになります。

    Compute Engine インスタンスに接続していない場合は、次のコマンドを実行して接続できます。

    TPU VM

    gcloud compute tpus tpu-vm ssh mnist-tutorial --zone=us-central1-b

    TPU ノード

    gcloud compute ssh mnist-tutorial --zone=us-central1-b

    これらの手順を続行する場合は、VM セッション ウィンドウで、(vm)$ で始まる各コマンドを実行します。

  8. TPU 名の環境変数を作成します。

    TPU VM

    (vm)$ export TPU_NAME=local

    TPU ノード

    (vm)$ export TPU_NAME=mnist-tutorial

  9. TensorFlow の要件をインストールします。

    使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。

    TPU VM

    (vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt

    TPU ノード

    (vm)$ pip3 install --user tensorflow-model-optimization>=0.1.3

モデルのトレーニング

MNIST TPU モデルのソースコードは GitHub で入手できます。

  1. 以下の変数を設定します。bucket-name は、使用するバケット名に置き換えてください。

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
(vm)$ export MODEL_DIR=${STORAGE_BUCKET}/mnist
(vm)$ export DATA_DIR=${STORAGE_BUCKET}/data

  2. PYTHONPATH 環境変数を設定する。

    TPU VM

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/tpu/models"

    TPU ノード

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"

  3. TPU を作成するときに、--version パラメータを -pjrt で終わるバージョンに設定した場合は、次の環境変数を設定して PJRT ランタイムを有効にします。

      (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
  (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

  4. モデルが保存されているディレクトリに移動します。

    TPU VM

    (vm)$ cd /usr/share/tpu/models/official/legacy/image_classification

    TPU ノード

    (vm)$ cd /usr/share/models/official/legacy/image_classification

  5. MNIST トレーニング スクリプトを実行します。

    (vm)$ python3 mnist_main.py \
  --tpu=${TPU_NAME} \
  --model_dir=${MODEL_DIR} \
  --data_dir=${DATA_DIR} \
  --train_epochs=10 \
  --distribution_strategy=tpu \
  --download

    コマンドフラグの説明

    tpu
    Cloud TPU の名前。Compute Engine VM と Cloud TPU の設定時に名前を指定しなかった場合は、デフォルトでユーザー名が指定されます。
    model_dir
    トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に生成されたチェックポイントを読み込むことができます。
    data_dir
    トレーニング入力の Cloud Storage パス。この例では、fake_imagenet データセットに設定されています。
    train_epochs
    モデルをトレーニングするためのエポックの数。
    distribution_strategy
    Cloud TPU で ResNet モデルをトレーニングするには、distribution_strategy を tpu に設定します。
    download
    true に設定すると、スクリプトが MNIST データセットをダウンロードして前処理します（ダウンロードがまだの場合）。

トレーニング スクリプトが v3-8 Cloud TPU で 5 分未満で実行され、次のような出力が表示されます。

Run stats:
{
  'accuracy_top_1': 0.9762369990348816,
  'eval_loss': 0.07863274961709976,
  'loss': 0.1111728847026825,
  'training_accuracy_top_1': 0.966645359992981
}

クリーンアップ

このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

  1. Compute Engine インスタンスとの接続を切断していない場合は切断します。

    (vm)$ exit

    プロンプトが username@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。

  2. Cloud TPU と Compute Engine リソースを削除します。リソースの削除に使用するコマンドは、TPU VM または TPU ノードのどちらを使用するかによって異なります。詳細については、システム アーキテクチャをご覧ください。

    TPU VM

    $ gcloud compute tpus tpu-vm delete mnist-tutorial \
--zone=us-central1-b

    TPU ノード

    $ gcloud compute tpus execution-groups delete mnist-tutorial \
--zone=us-central1-b

  3. gcloud compute tpus execution-groups list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。次のコマンドの出力には、このチュートリアルで作成したリソースを含めないでください。

    TPU VM

    $ gcloud compute tpus tpu-vm list --zone=us-central1-b

    TPU ノード

    $ gcloud compute tpus execution-groups list --zone=us-central1-b

  4. 次に示すように、gsutil を使用して Cloud Storage バケットを削除します。bucket-name を Cloud Storage バケットの名前に置き換えます。

    $ gsutil rm -r gs://bucket-name

次のステップ

TensorFlow Cloud TPU のチュートリアルでは通常、サンプル データセットを使用してモデルをトレーニングします。このトレーニングの結果は推論に使用できません。モデルを推論に使用するには、一般公開されているデータセットまたは独自のデータセットでデータをトレーニングします。Cloud TPU でトレーニングされた TensorFlow モデルは通常、データセットを TFRecord 形式にする必要があります。

データセット変換ツールのサンプルを使用して、画像分類データセットを TFRecord 形式に変換できます。画像分類モデルを使用しない場合は、自分でデータセットを TFRecord 形式に変換する必要があります。詳細については、TFRecord と tf.Example をご覧ください。

ハイパーパラメータ チューニング

データセットでモデルのパフォーマンスを向上させるには、モデルのハイパーパラメータを調整します。すべての TPU でサポートされているモデルに共通のハイパーパラメータに関する情報については、GitHub をご覧ください。モデルに固有のハイパーパラメータに関する情報については、各モデルのソースコードで確認できます。ハイパーパラメータ調整の詳細については、ハイパーパラメータ調整の概要とハイパーパラメータを調整するをご覧ください。

推論

モデルをトレーニングしたら、そのモデルを推論（予測）に使用できます。Cloud TPU 推論コンバータ ツールを使用して、Cloud TPU v5e での推論用の TensorFlow モデルを準備して最適化できます。Cloud TPU v5e での推論の詳細については、Cloud TPU v5e 推論の概要をご覧ください。