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 で ML モデルを作成して実行する場合に推奨される方法です。

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 バケットを作成します。

    gcloud storage buckets create gs://bucket-name \
--project=${PROJECT_ID} \
--location=us-central1

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

  6. gcloud コマンドを使用して TPU VM を作成します。

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

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン。

    accelerator-type
    アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。

    version
    Cloud TPU のソフトウェア バージョン。

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

  7. SSH を使用して Compute Engine インスタンスに接続します。VM に接続すると、シェル プロンプトが username@projectname から username@vm-name に変わります。

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

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

    (vm)$ export TPU_NAME=local

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

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

モデルのトレーニング

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 環境変数を設定する。

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/tpu/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. モデルが保存されているディレクトリに移動します。

    (vm)$ cd /usr/share/tpu/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 リソースを削除します。

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

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

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

  4. 次のコマンドを使用して Cloud Storage バケットを削除します。bucket-name は、使用する Cloud Storage バケットの名前に置き換えます。

    $ gcloud storage rm gs://bucket-name --recursive

次のステップ

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 推論の概要をご覧ください。