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 を作成する
gcloud
コマンドを使用して Compute Engine VM と Cloud TPU を起動します。使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。2 つの VM アーキテクチャの詳細については、システム アーキテクチャをご覧ください。gcloud
コマンドの詳細については、gcloud リファレンスをご覧ください。
TPU VM
デフォルトのユーザー プロジェクトで TPU VM を作成するには、ネットワークと compute/zone を実行します。
$ gcloud compute tpus tpu-vm create tpu-name \
--zone=us-central1-b \
--accelerator-type=v3-8 \
--version=tpu-vm-pt-1.11
TPU を作成するときに、デフォルトのネットワークとサブネットワークを指定する場合は、追加の --network
フラグと --subnetwork
フラグを渡します。デフォルトのネットワークを使用しない場合は、--network
フラグを渡す必要があります。--subnetwork
フラグは省略可能で、使用するネットワーク(デフォルトまたはユーザー指定)のデフォルトのサブネットワークを指定するために使用します。これらのフラグの詳細については、gcloud
API リファレンス ページをご覧ください。
TPU ノード
PyTorch 用の TPU ノードを作成する場合、まず Compute Engine VM インスタンスを作成します。
gcloud compute instances create tpu-name \ --zone=us-central1-b \ --machine-type=n1-standard-16 \ --image-family=torch-xla \ --image-project=ml-images \ --boot-disk-size=200GB \ --scopes=https://www.googleapis.com/auth/cloud-platform
次に、TPU インスタンスを作成します。
gcloud compute tpus create tpu-name \ --zone=us-central1-b \ --network=default \ --version=pytorch-1.11 \ --accelerator-type=v3-8
Cloud TPU VM に接続する
TPU VM
$ gcloud compute tpus tpu-vm ssh tpu-name \
--zone us-central1-b
TPU ノード
gcloud compute ssh tpu-name --zone=us-central1-b
XRT TPU デバイス構成を設定する
TPU VM
Torch-XLA 環境を構成します。
(vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"
TPU ノード
TPU ノードの IP アドレスを調べます。
(vm)$ gcloud compute tpus describe tpu-name --zone=us-central1-b
Torch-XLA 環境を構成します。your-tpu-ip-address は、TPU の IP アドレスに置き換えてください。
(vm)$ conda activate torch-xla-1.11 (vm)$ export TPU_IP_ADDRESS=your-tpu-ip-address (vm)$ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
簡単な計算を実行する
現在のディレクトリに
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')
OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op: TPURoundRobin OpKernel ('op: "TpuHandleToProtoKey" device_type: "CPU"') for unknown op: TpuHandleToProtoKey
クリーンアップ
このページで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、次の手順を行います。
Compute Engine インスタンスとの接続を切断していない場合は切断します。
(vm)$ exit
プロンプトが
username@projectname
に変わります。これは、現在、Cloud Shell 内にいることを示しています。Cloud TPU を削除します。
TPU VM
$ gcloud compute tpus tpu-vm delete tpu-name \ --zone=us-central1-b
TPU ノード
$ gcloud compute tpus execution-groups delete tpu-name \ --zone=us-central1-b
このコマンドの出力で、TPU が削除されたことがわかります。
次のステップ
Cloud TPU VM の詳細を確認する。