PyTorch を使用して Cloud TPU VM で計算を実行する

このクイックスタートでは、Cloud TPU の作成、PyTorch のインストール、Cloud TPU で簡単な計算を行う方法を説明します。Cloud TPU でモデルをトレーニングする方法の詳細については、Cloud TPU PyTorch のチュートリアルをご覧ください。

準備

このクイックスタートを使用する前に、Google Cloud Platform アカウントを作成し、Google Cloud CLI をインストールして gcloud コマンドを構成する必要があります。詳細しくは、アカウントと Cloud TPU プロジェクトを設定するをご覧ください。

`gcloud` を使用して Cloud TPU を作成する

デフォルトのユーザープロジェクトで TPU VM を作成するには、ネットワークと zone を実行します。

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central1-a \
   --accelerator-type=v3-8 \
   --version=tpu-ubuntu2204-base

コマンドフラグの説明

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

TPU を作成するときに、デフォルトのネットワークとサブネットワークを指定する場合は、追加の --network フラグと --subnetwork フラグを渡します。デフォルトネットワークを使用しない場合は、--network フラグを渡す必要があります。--subnetwork フラグは省略可能で、使用するネットワーク（デフォルトまたはユーザー指定）のデフォルトのサブネットワークを指定するために使用します。これらのフラグの詳細については、gcloud API リファレンスページをご覧ください。

Cloud TPU VM に接続する

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-a

TPU VM に PyTorch/XLA をインストールする

   (vm)$ pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html

TPU ランタイム構成を設定する

PyTorch/XLA ランタイムが TPU を使用していることを確認します。

   (vm) $ export PJRT_DEVICE=TPU

PyTorch が TPU にアクセスできることを確認する

現在のディレクトリに tpu-count.py という名前のファイルを作成し、次のスクリプトをコピーして貼り付けます。
```
import torch
import torch_xla
import torch_xla.core.xla_model as xm
print(f'PyTorch can access {len(torch_xla.devices())} TPU cores')
```
スクリプトを実行します。
```
(vm)$ python3 tpu-count.py
```
スクリプトからの出力には、計算の結果が表示されます。
```
PyTorch can access 8 TPU cores
```

基本的な計算を行う

現在のディレクトリに tpu-test.py という名前のファイルを作成し、次のスクリプトをコピーして貼り付けます。

import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

スクリプトを実行します。

  (vm)$ python3 tpu-test.py

スクリプトからの出力には、計算の結果が表示されます。

tensor([[-0.2121,  1.5589, -0.6951],
        [-0.7886, -0.2022,  0.9242],
        [ 0.8555, -1.8698,  1.4333]], device='xla:1')

クリーンアップ

このページで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、次の手順を行います。

Compute Engine インスタンスとの接続を切断していない場合は切断します。
```
(vm)$ exit
```
プロンプトが username@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。

Cloud TPU を削除します。

$ gcloud compute tpus tpu-vm delete tpu-name \
  --zone=us-central1-a

このコマンドの出力で、TPU が削除されたことがわかります。

次のステップ

Cloud TPU VM の詳細を確認する。