Docker コンテナで Cloud TPU アプリケーションを実行する
Docker コンテナでは、コードと必要なすべての依存関係を 1 つの配布可能なパッケージにまとめることで、アプリケーションの構成が容易になります。TPU VM 内で Docker コンテナを実行して、Cloud TPU アプリケーションの構成と共有を簡素化できます。このドキュメントでは、Cloud TPU でサポートされている機械学習フレームワークごとに Docker コンテナを設定する方法について説明します。
Docker コンテナで TensorFlow モデルをトレーニングする
TPU デバイス
現在のディレクトリに
Dockerfile
という名前のファイルを作成し、次のテキストを貼り付けます。FROM python:3.8 RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.10.0/tensorflow-2.10.0-cp38-cp38-linux_x86_64.whl RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.4.0/libtpu.so -o /lib/libtpu.so RUN git clone https://github.com/tensorflow/models.git WORKDIR ./models RUN pip install -r official/requirements.txt ENV PYTHONPATH=/models
Cloud Storage バケットを作成する
gsutil mb -c standard -l europe-west4 gs://your-bucket-name
TPU VM の作成
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-tf-2.11.0
Dockerfile を TPU VM にコピーする
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
TPU VM に SSH 接続する
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Docker イメージをビルドする
sudo docker build -t your-image-name .
Docker コンテナを起動する
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
環境変数を設定する
export STORAGE_BUCKET=gs://your-bucket-name export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x export ZONE=europe-west4-a export TPU_NAME=local
ResNet をトレーニングする
python3 official/vision/train.py \ --tpu=local \ --experiment=resnet_imagenet \ --mode=train_and_eval \ --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \ --model_dir=${MODEL_DIR} \ --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,trainer.train_steps=100"
TPU Pod
現在のディレクトリに
Dockerfile
という名前のファイルを作成し、次のテキストを貼り付けます。FROM python:3.8 RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.10.0/tensorflow-2.10.0-cp38-cp38-linux_x86_64.whl RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.4.0/libtpu.so -o /lib/libtpu.so RUN git clone https://github.com/tensorflow/models.git WORKDIR ./models RUN pip install -r official/requirements.txt ENV PYTHONPATH=/models
TPU VM の作成
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v3-32 \ --version=tpu-vm-tf-2.11.0-pod
Dockerfile を TPU VM にコピーする
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
TPU VM に SSH 接続する
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Docker イメージをビルドする
sudo docker build -t your-image-name .
Docker コンテナを起動する
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
ResNet をトレーニングする
python3 official/vision/train.py \ --tpu=your-tpu-name \ --experiment=resnet_imagenet \ --mode=train_and_eval \ --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \ --model_dir=${MODEL_DIR} \ --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
トレーニング スクリプトが完了したら、リソースをクリーンアップしてください。
- 「
exit
」と入力して Docker コンテナを終了します。 - 「
exit
」と入力して TPU VM を終了します。 - TPU VM を削除する
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
Docker コンテナで PyTorch モデルをトレーニングする
Cloud TPU VM を作成する
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-pt-1.13
TPU VM への SSH 接続
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
TPU VM でコンテナを起動する
sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm bash
libtpu のインストール
pip install torch_xla[tpuvm] export XRT_TPU_CONFIG="localservice;0;localhost:51011"
PyTorch XLA リポジトリのクローンを作成する
git clone --recursive https://github.com/pytorch/xla.git
ImageNet をトレーニングする
python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
トレーニング スクリプトが完了したら、リソースをクリーンアップしてください。
- 「
exit
」と入力して Docker コンテナを終了します。 - 「
exit
」と入力して TPU VM を終了します。 - TPU VM を削除する
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
Docker コンテナで JAX モデルをトレーニングする
TPU デバイス
TPU VM を作成します。
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-base
TPU VM への SSH 接続
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
TPU VM で Docker デーモンを起動する
sudo systemctl start docker
Docker コンテナを起動する
sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
JAX のインストール
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
FLAX のインストール
pip install --upgrade clu git clone https://github.com/google/flax.git pip install --user -e flax
Flax MNIST トレーニング スクリプトを実行します。
cd flax/examples/mnist python3 main.py --workdir=/tmp/mnist \ --config=configs/default.py \ --config.learning_rate=0.05 \ --config.num_epochs=5
トレーニング スクリプトが完了したら、リソースをクリーンアップしてください。
- 「
exit
」と入力して Docker コンテナを終了します。 - 「
exit
」と入力して TPU VM を終了します。 TPU VM の削除
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
TPU Pod
TPU Pod で JAX コードを実行するときは、すべての TPU ワーカーで同時に JAX コードを実行する必要があります。これを行う方法の 1 つとして、--worker=all
フラグと --command
フラグを指定した gcloud compute tpus tpu-vm ssh
コマンドを使用する方法があります。次の手順は、各 Docker ワーカーの設定を容易にする Docker イメージの作成方法を示しています。
ターミナル プロンプトで、次の環境変数を設定します。
export USER=your-user-id export PROJECT=your-project-name export REPO_NAME=your-repo-name export IMAGE_NAME=your-image-name export TAG=your-tag export TPU_NAME=your-tpu-name export ZONE=europe-west4-a export ACCEL_TYPE=v2-8
次のように置き換えます。
your-user-id
は、実際のユーザー ID に置き換えます。your-project-name
は、Google Cloud プロジェクト名に置き換えます。your-repo-name
は、Artifact Registry リポジトリの名前に置き換えます。your-image-name
は、Docker イメージの名前に置き換えます。your-tag
は、Docker イメージのタグに置き換えます。your-tpu-name
は、TPU の名前に置き換えます。
現在のディレクトリに
Dockerfile
という名前のファイルを作成し、次のテキストを貼り付けます。FROM google-cloud-cli:latest RUN pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html RUN pip install --upgrade clu RUN git clone https://github.com/google/flax.git RUN pip install --user -e flax WORKDIR ./flax/examples/mnist
Docker イメージをビルドする
docker build -t $IMAGE .
Artifact Registry に push する前に、Docker イメージにタグを追加します。Artifact Registry の操作の詳細については、コンテナ イメージの操作をご覧ください。
docker tag $IMAGE_NAME europe-west-docker.pkg.dev/$PROJECT_NAME/$REPO_NAME/$IMAGE_NAME:$TAG
Docker イメージを Artifact Registry に push する
docker push europe-west4-docker.pkg.dev/ml-writers/$REPO_NAME/$IMAGE_NAME:$TAG
TPU VM の作成
gcloud compute tpus tpu-vm create $TPU_NAME \ --zone $ZONE \ --accelerator-type $ACCEL_TYPE \ --version tpu-vm-base
すべての TPU ワーカーで、Artifact Registry から Docker イメージを pull します。
gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --command="sudo usermod -a -G docker $USER" gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet" gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --command="docker pull europe-west4-docker.pkg.dev/your-project-name/your-repo-name/your-image-name:your-tag"
すべての TPU ワーカーでコンテナを実行します。
gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project-name/your-repo-name/your-image-name:your-tag bash"
すべての TPU ワーカーでトレーニング スクリプトを実行します。
gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --zone $ZONE \ --command="docker exec --privileged $CONTAINER_NAME python3 main.py --workdir=/tmp/mnist \ --config=configs/default.py \ --config.learning_rate=0.05 \ --config.num_epochs=5"
トレーニング スクリプトが完了したら、リソースをクリーンアップしてください。次のコマンドを使用して TPU VM を削除します。
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
次のステップ
- Cloud TPU のチュートリアル
- Cloud TPU ユーザーガイド
- Cloud TPU システム アーキテクチャ
- TPU Pod スライスでの TensorFlow コードの実行
- TPU Pod スライスでの JAX コードの実行