GKE Autopilot モードで GPU を使用してモデルをトレーニングする

サンプル リポジトリのクローンを作成する

Cloud Shell で、次のコマンドを実行します。

git clone https://github.com/GoogleCloudPlatform/ai-on-gke && \
cd ai-on-gke/tutorials-and-examples/gpu-examples/training-single-gpu

クラスタを作成する

1. Google Cloud コンソールで、[Autopilot クラスタの作成] ページに移動します。

   [Autopilot クラスタの作成] に移動

2. [名前] フィールドに「gke-gpu-cluster」と入力します。

3. [リージョン] リストで [us-central1] を選択します。

4. [作成] をクリックします。

Cloud Storage バケットを作成する

1. Google Cloud コンソールで、[バケットの作成] ページに移動します。

   [バケットの作成] に移動

2. [バケットに名前を付ける] フィールドに、次の名前を入力します。

   PROJECT_ID-gke-gpu-bucket
   

   PROJECT_ID は、実際の Google Cloud プロジェクト ID に置き換えます。

3. [続行] をクリックします。

4. [ロケーション タイプ] で、[リージョン] を選択します。

5. [リージョン] リストで us-central1 (Iowa) を選択し、[続行] をクリックします。

6. [データのストレージ クラスを選択する] セクションで、[続行] をクリックします。

7. [オブジェクトへのアクセスを制御する方法を選択する] セクションで、[アクセス制御] には [均一] を選択します。

8. [作成] をクリックします。

9. [公開アクセスの防止] ダイアログで、[このバケットに対する公開アクセス禁止を適用する] チェックボックスをオンにして、[確認] をクリックします。

Workload Identity Federation for GKE を使用してバケットにアクセスするようにクラスタを構成する

クラスタから Cloud Storage バケットにアクセスできるようにする手順は次のとおりです。

1. クラスタに Kubernetes ServiceAccount を作成する。
2. ServiceAccount がバケットにアクセスできるように IAM 許可ポリシーを作成する。

クラスタに Kubernetes ServiceAccount を作成する

Cloud Shell で、次の操作を行います。

1. クラスタに接続します。

   gcloud container clusters get-credentials gke-gpu-cluster \
       --location=us-central1
   

2. Kubernetes Namespace を作成します。

   kubectl create namespace gke-gpu-namespace
   

3. Namespace に Kubernetes ServiceAccount を作成します。

   kubectl create serviceaccount gpu-k8s-sa --namespace=gke-gpu-namespace
   

バケットに IAM 許可ポリシーを作成する

バケットに対する Storage オブジェクト管理者（roles/storage.objectAdmin）ロールを Kubernetes ServiceAccount に付与します。

gcloud storage buckets add-iam-policy-binding gs://PROJECT_ID \
    --member=principal://iam.googleapis.com/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/gke-gpu-namespace/sa/gpu-k8s-sa \
    --role=roles/storage.objectAdmin \
    --condition=None

PROJECT_NUMBER は、実際の Google Cloud プロジェクトの番号に置き換えます。

Pod が Cloud Storage バケットにアクセスできることを確認する

1. Cloud Shell で、次の環境変数を作成します。

   export K8S_SA_NAME=gpu-k8s-sa
   export BUCKET_NAME=PROJECT_ID-gke-gpu-bucket
   

   PROJECT_ID は、実際の Google Cloud プロジェクト ID に置き換えます。

2. TensorFlow コンテナを含む Pod を作成します。

   envsubst < src/gke-config/standard-tensorflow-bash.yaml | kubectl --namespace=gke-gpu-namespace apply -f -
   

   このコマンドで、作成した環境変数がマニフェスト内の対応する参照に挿入されます。マニフェストをテキスト エディタで開いて、$K8S_SA_NAME と $BUCKET_NAME を対応する値に置き換えることもできます。

3. バケットにサンプル ファイルを作成します。

   touch sample-file
   gsutil cp sample-file gs://PROJECT_ID-gke-gpu-bucket
   

4. Pod の準備ができるまで待ちます。

   kubectl wait --for=condition=Ready pod/test-tensorflow-pod -n=gke-gpu-namespace --timeout=180s
   

   Pod の準備ができると、出力は次のようになります。

   pod/test-tensorflow-pod condition met
   

   コマンドがタイムアウトした場合、GKE は Pod を実行する新しいノードをまだ作成している可能性があります。コマンドをもう一度実行し、Pod の準備が整うまで待ちます。

5. TensorFlow コンテナでシェルを開きます。

   kubectl -n gke-gpu-namespace exec --stdin --tty test-tensorflow-pod --container tensorflow -- /bin/bash
   

6. 作成したサンプル ファイルを閲覧してみます。

   ls /data
   

   出力にサンプル ファイルが表示されます。

7. ログを確認して Pod に接続されている GPU を特定します。

   python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"
   

   出力には、Pod に接続された GPU が次のように表示されます。

   ...
   PhysicalDevice(name='/physical_device:GPU:0',device_type='GPU')
   

8. コンテナを終了します。

   exit
   

9. サンプル Pod を削除します。

   kubectl delete -f src/gke-config/standard-tensorflow-bash.yaml \
       --namespace=gke-gpu-namespace
   

MNIST データセットを使用してトレーニングと予測を行う

このセクションでは、サンプル データセット MNIST でトレーニング ワークロードを実行します。

1. サンプルデータを Cloud Storage バケットにコピーします。

   gsutil -m cp -R src/tensorflow-mnist-example gs://PROJECT_ID-gke-gpu-bucket/
   

2. 次の環境変数を作成します。

   export K8S_SA_NAME=gpu-k8s-sa
   export BUCKET_NAME=PROJECT_ID-gke-gpu-bucket
   

3. トレーニング Job を確認します。

   # Copyright 2023 Google LLC
   #
   # Licensed under the Apache License, Version 2.0 (the "License");
   # you may not use this file except in compliance with the License.
   # You may obtain a copy of the License at
   #
   #      http://www.apache.org/licenses/LICENSE-2.0
   #
   # Unless required by applicable law or agreed to in writing, software
   # distributed under the License is distributed on an "AS IS" BASIS,
   # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   # See the License for the specific language governing permissions and
   # limitations under the License.
   
   apiVersion: batch/v1
   kind: Job
   metadata:
     name: mnist-training-job
   spec:
     template:
       metadata:
         name: mnist
         annotations:
           gke-gcsfuse/volumes: "true"
       spec:
         nodeSelector:
           cloud.google.com/gke-accelerator: nvidia-tesla-t4
         tolerations:
         - key: "nvidia.com/gpu"
           operator: "Exists"
           effect: "NoSchedule"
         containers:
         - name: tensorflow
           image: tensorflow/tensorflow:latest-gpu 
           command: ["/bin/bash", "-c", "--"]
           args: ["cd /data/tensorflow-mnist-example; pip install -r requirements.txt; python tensorflow_mnist_train_distributed.py"]
           resources:
             limits:
               nvidia.com/gpu: 1
               cpu: 1
               memory: 3Gi
           volumeMounts:
           - name: gcs-fuse-csi-vol
             mountPath: /data
             readOnly: false
         serviceAccountName: $K8S_SA_NAME
         volumes:
         - name: gcs-fuse-csi-vol
           csi:
             driver: gcsfuse.csi.storage.gke.io
             readOnly: false
             volumeAttributes:
               bucketName: $BUCKET_NAME
               mountOptions: "implicit-dirs"
         restartPolicy: "Never"

4. トレーニング Job をデプロイします。

   envsubst < src/gke-config/standard-tf-mnist-train.yaml | kubectl -n gke-gpu-namespace apply -f -
   

   このコマンドで、作成した環境変数がマニフェスト内の対応する参照に置き換えます。マニフェストをテキスト エディタで開いて、$K8S_SA_NAME と $BUCKET_NAME を対応する値に置き換えることもできます。

5. Job のステータスが Completed になるまで待ちます。

   kubectl wait -n gke-gpu-namespace --for=condition=Complete job/mnist-training-job --timeout=180s
   

   Job の準備ができている場合、出力は次のようになります。

   job.batch/mnist-training-job condition met
   

   コマンドがタイムアウトした場合、GKE は Pod を実行する新しいノードをまだ作成している可能性があります。コマンドをもう一度実行し、Job の準備が整うまで待ちます。

6. TensorFlow コンテナのログを確認します。

   kubectl logs -f jobs/mnist-training-job -c tensorflow -n gke-gpu-namespace
   

   出力には、以下のイベントの発生が示されます。

   • 必要な Python パッケージのインストール
   • MNIST データセットのダウンロード
   • GPU を使用したモデルのトレーニング
   • モデルを保存する
   • モデルの評価

   ...
   Epoch 12/12
   927/938 [============================>.] - ETA: 0s - loss: 0.0188 - accuracy: 0.9954
   Learning rate for epoch 12 is 9.999999747378752e-06
   938/938 [==============================] - 5s 6ms/step - loss: 0.0187 - accuracy: 0.9954 - lr: 1.0000e-05
   157/157 [==============================] - 1s 4ms/step - loss: 0.0424 - accuracy: 0.9861
   Eval loss: 0.04236088693141937, Eval accuracy: 0.9861000180244446
   Training finished. Model saved
   

7. トレーニング ワークロードを削除します。

   kubectl -n gke-gpu-namespace delete -f src/gke-config/standard-tf-mnist-train.yaml
   

推論ワークロードをデプロイする

このセクションでは、サンプル データセットを入力として受け取り、予測を返す推論ワークロードをデプロイします。

1. 予測に使用する画像をバケットにコピーします。

   gsutil -m cp -R data/mnist_predict gs://PROJECT_ID-gke-gpu-bucket/
   

2. 推論ワークロードを確認します。

   # Copyright 2023 Google LLC
   #
   # Licensed under the Apache License, Version 2.0 (the "License");
   # you may not use this file except in compliance with the License.
   # You may obtain a copy of the License at
   #
   #      http://www.apache.org/licenses/LICENSE-2.0
   #
   # Unless required by applicable law or agreed to in writing, software
   # distributed under the License is distributed on an "AS IS" BASIS,
   # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   # See the License for the specific language governing permissions and
   # limitations under the License.
   
   apiVersion: batch/v1
   kind: Job
   metadata:
     name: mnist-batch-prediction-job
   spec:
     template:
       metadata:
         name: mnist
         annotations:
           gke-gcsfuse/volumes: "true"
       spec:
         nodeSelector:
           cloud.google.com/gke-accelerator: nvidia-tesla-t4
         tolerations:
         - key: "nvidia.com/gpu"
           operator: "Exists"
           effect: "NoSchedule"
         containers:
         - name: tensorflow
           image: tensorflow/tensorflow:latest-gpu 
           command: ["/bin/bash", "-c", "--"]
           args: ["cd /data/tensorflow-mnist-example; pip install -r requirements.txt; python tensorflow_mnist_batch_predict.py"]
           resources:
             limits:
               nvidia.com/gpu: 1
               cpu: 1
               memory: 3Gi
           volumeMounts:
           - name: gcs-fuse-csi-vol
             mountPath: /data
             readOnly: false
         serviceAccountName: $K8S_SA_NAME
         volumes:
         - name: gcs-fuse-csi-vol
           csi:
             driver: gcsfuse.csi.storage.gke.io
             readOnly: false
             volumeAttributes:
               bucketName: $BUCKET_NAME
               mountOptions: "implicit-dirs"
         restartPolicy: "Never"

3. 推論ワークロードをデプロイします。

   envsubst < src/gke-config/standard-tf-mnist-batch-predict.yaml | kubectl -n gke-gpu-namespace apply -f -
   

   このコマンドで、作成した環境変数がマニフェスト内の対応する参照に置き換えます。マニフェストをテキスト エディタで開いて、$K8S_SA_NAME と $BUCKET_NAME を対応する値に置き換えることもできます。

4. Job のステータスが Completed になるまで待ちます。

   kubectl wait -n gke-gpu-namespace --for=condition=Complete job/mnist-batch-prediction-job --timeout=180s
   

   出力は次のようになります。

   job.batch/mnist-batch-prediction-job condition met
   

5. TensorFlow コンテナのログを確認します。

   kubectl logs -f jobs/mnist-batch-prediction-job -c tensorflow -n gke-gpu-namespace
   

   出力は、次のような各画像の予測とモデルの信頼度です。

   Found 10 files belonging to 1 classes.
   1/1 [==============================] - 2s 2s/step
   The image /data/mnist_predict/0.png is the number 0 with a 100.00 percent confidence.
   The image /data/mnist_predict/1.png is the number 1 with a 99.99 percent confidence.
   The image /data/mnist_predict/2.png is the number 2 with a 100.00 percent confidence.
   The image /data/mnist_predict/3.png is the number 3 with a 99.95 percent confidence.
   The image /data/mnist_predict/4.png is the number 4 with a 100.00 percent confidence.
   The image /data/mnist_predict/5.png is the number 5 with a 100.00 percent confidence.
   The image /data/mnist_predict/6.png is the number 6 with a 99.97 percent confidence.
   The image /data/mnist_predict/7.png is the number 7 with a 100.00 percent confidence.
   The image /data/mnist_predict/8.png is the number 8 with a 100.00 percent confidence.
   The image /data/mnist_predict/9.png is the number 9 with a 99.65 percent confidence.