PyTorch を使用した拡散モデルのトレーニング

---

このチュートリアルでは、PyTorch Lightning と Pytorch XLA を使用して TPU で拡散モデルをトレーニングする方法について説明します。

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

目標

• Cloud TPU の作成
• PyTorch Lightning をインストールする
• 拡散リポジトリのクローンを作成する
• Imagenette データセットを準備する
• トレーニング スクリプトを実行します

費用

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

• Compute Engine
• Cloud TPU

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

始める前に

このチュートリアルを開始する前に、Google Cloud プロジェクトが正しく設定されていることを確認します。

1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

3. Make sure that billing is enabled for your Google Cloud project.

4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

5. Make sure that billing is enabled for your Google Cloud project.

このチュートリアルでは、Google Cloud の課金対象となるコンポーネントを使用します。費用を見積もるには、Cloud TPU の料金ページを確認してください。不要な課金を回避するために、このチュートリアルを完了したら、作成したリソースを必ずクリーンアップしてください。

Cloud TPU の作成

このチュートリアルでは v4-8 を使用しますが、単一ホストのすべてのアクセラレータ サイズで同様に動作します。

コマンドを簡単に使用できるように、いくつかの環境変数を設定します。

export ZONE=us-central2-b
export PROJECT_ID=your-project-id
export ACCELERATOR_TYPE=v4-8
export RUNTIME_VERSION=tpu-ubuntu2204-base
export TPU_NAME=your_tpu_name

Cloud TPU の作成

注: 現在、TPU Pod を作成するのに十分な容量がない場合は、キューに入れられたリソースを使用して、リクエストをキューに入れることができます。キューに入れられたリソースを使用すると、容量が利用可能になり次第、容量を受け取ることができます。Cloud TPU リソースをキューに入れられたリソースとしてリクエストするには、代わりに gcloud compute tpus queued-resources create コマンドを使用します。詳細については、キューに入れられたリソースを管理するをご覧ください。

gcloud compute tpus tpu-vm create ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--accelerator-type=${ACCELERATOR_TYPE} \
--version=${RUNTIME_VERSION} \
--subnetwork=tpusubnet

注: 次のコマンドは、すべての TPU VM で同時に実行されます。

必要なソフトウェアのインストール

1. PyTorch / XLA の最新リリース v2.4.0 とともに、必要なパッケージをインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --command="sudo apt-get update -y && sudo apt-get install libgl1 -y
   git clone https://github.com/pytorch-tpu/stable-diffusion.git
   cd stable-diffusion
   pip install -r requirements.txt
   pip install -e .
   pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/master.zip -U
   pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
   pip install clip
   pip install torch~=2.4.0 torch_xla[tpu]~=2.4.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"

2. torch 2.2 以降と互換性があるようにソースファイルを修正します。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --command="cd stable-diffusion/
   sed -i 's/from torch._six import string_classes/string_classes = (str, bytes)/g' src/taming-transformers/taming/data/utils.py
   sed -i 's/trainer_kwargs\\[\"callbacks\"\\]/# trainer_kwargs\\[\"callbacks\"\\]/g' main_tpu.py"

3. Imagenette（Imagenet データセットの小さいバージョン）をダウンロードして、適切なディレクトリに移動します。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --command="wget -nv https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz
   tar -xf  imagenette2.tgz
   mkdir -p ~/.cache/autoencoders/data/ILSVRC2012_train/data
   mkdir -p ~/.cache/autoencoders/data/ILSVRC2012_validation/data
   mv imagenette2/train/*  ~/.cache/autoencoders/data/ILSVRC2012_train/data
   mv imagenette2/val/* ~/.cache/autoencoders/data/ILSVRC2012_validation/data"

4. 最初のステージの事前トレーニング済みモデルをダウンロードします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --command="cd stable-diffusion/
   wget -nv -O models/first_stage_models/vq-f8/model.zip https://ommer-lab.com/files/latent-diffusion/vq-f8.zip
   cd  models/first_stage_models/vq-f8/
   unzip -o model.zip"

モデルのトレーニング

次のコマンドでトレーニングを実行します。v4-8 では、トレーニング プロセスに約 30 分かかることが想定されます。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--command="python3 stable-diffusion/main_tpu.py --train --no-test --base=stable-diffusion/configs/latent-diffusion/cin-ldm-vq-f8-ss.yaml -- data.params.batch_size=32 lightning.trainer.max_epochs=5 model.params.first_stage_config.params.ckpt_path=stable-diffusion/models/first_stage_models/vq-f8/model.ckpt lightning.trainer.enable_checkpointing=False lightning.strategy.sync_module_states=False"

クリーンアップ

作成したリソースを使用した後、アカウントに不要な請求が発生しないようにクリーンアップを行います。

Google Cloud CLI を使用して Cloud TPU リソースを削除します。

  $  gcloud compute tpus tpu-vm delete diffusion-tutorial --zone=us-central2-b
  

次のステップ

• PyTorch を使用した Cloud TPU での Resnet50 のトレーニング
• TPU での Pytorch のトラブルシューティング
• PyTorch / XLA のドキュメント