打开一个 Cloud Shell 窗口。

打开 Cloud Shell
为项目 ID 创建一个变量。
```
export PROJECT_ID=project-id
```
配置 Google Cloud CLI，以使用您要创建 Cloud TPU 的项目。
```
gcloud config set project ${PROJECT_ID}
```
当您第一次在新的 Cloud Shell 虚拟机中运行此命令时，系统会显示 Authorize Cloud Shell 页面。点击页面底部的 Authorize 以允许 gcloud 使用您的凭据进行 Google Cloud API 调用。

从 Cloud Shell 启动本教程所需的 Compute Engine 资源。注意：如果您在 Criteo Kaggle 数据集上进行训练，则需要使用 n1-highmem-96 machine-type。

gcloud compute instances create dlrm-tutorial \
--zone=us-central1-a \
--machine-type=n1-standard-64 \
--image-family=torch-xla \
--image-project=ml-images  \
--boot-disk-size=200GB \
--scopes=https://www.googleapis.com/auth/cloud-platform

连接到新的 Compute Engine 实例。
```
gcloud compute ssh dlrm-tutorial --zone=us-central1-a
```
要点：从现在起，前缀 (vm) $ 表示您应在 Compute Engine 虚拟机实例上运行该命令。

启动 Cloud TPU 资源

在 Compute Engine 虚拟机中，使用以下命令启动 Cloud TPU 资源：

(vm) $ gcloud compute tpus create dlrm-tutorial \
--zone=us-central1-a \
--network=default \
--version=pytorch-2.0  \
--accelerator-type=v3-8

确定 Cloud TPU 资源的 IP 地址。
```
(vm) $ gcloud compute tpus describe dlrm-tutorial --zone=us-central1-a
```
重要提示 ：内部 IP 地址也会显示在 Google Cloud 控制台的 Compute Engine > TPU 下。在创建和配置 PyTorch 环境时，您需要不带端口号的 IP 地址。

创建并配置 PyTorch 环境

启动 conda 环境。
```
(vm) $ conda activate torch-xla-2.0
```
为 Cloud TPU 资源配置环境变量。

注意：TPU_IP_ADDRESS 变量必须与您在启动 Cloud TPU 资源时确定的 Cloud TPU 的 IP 地址相等。
```
(vm) $ export TPU_IP_ADDRESS=ip-address
```
```
(vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
```

使用虚构数据运行训练作业

安装依赖项。
```
(vm) $ pip install onnx
```

对随机数据运行模型。此过程将需要 5-10 分钟。

(vm) $ python /usr/share/torch-xla-2.0/tpu-examples/deps/dlrm/dlrm_tpu_runner.py \
    --arch-embedding-size=1000000-1000000-1000000-1000000-1000000-1000000-1000000-1000000 \
    --arch-sparse-feature-size=64 \
    --arch-mlp-bot=512-512-64 \
    --arch-mlp-top=1024-1024-1024-1 \
    --arch-interaction-op=dot \
    --lr-num-warmup-steps=10 \
    --lr-decay-start-step=10 \
    --mini-batch-size=2048 \
    --num-batches=1000 \
    --data-generation='random' \
    --numpy-rand-seed=727 \
    --print-time \
    --print-freq=100 \
    --num-indices-per-lookup=100 \
    --use-tpu \
    --num-indices-per-lookup-fixed \
    --tpu-model-parallel-group-len=8 \
    --tpu-metrics-debug \
    --tpu-cores=8

（可选）在 Criteo Kaggle 数据集上训练

以下是可选步骤。仅当您希望在 Criteo Kaggle 数据集上训练时才应运行这些步骤。

下载数据集。

按照此处的说明从 Criteo Kaggle 数据集中下载数据集。下载完成后，将 dac.tar.gz 文件复制到名为 ./criteo-kaggle/ 的目录中。使用 tar -xzvf 命令将 tar.gz 文件的内容解压缩到 ./critero-kaggle 目录中。
```
 (vm) $ mkdir criteo-kaggle
 (vm) $ cd criteo-kaggle
 (vm) $ # Download dataset from above link here.
 (vm) $ tar -xzvf dac.tar.gz
 (vm) $ cd ..
```

预处理数据集。

启动此脚本以预处理 Criteo 数据集。此脚本会生成一个名为 kaggleAdDisplayChallenge_processed.npz 的文件，并且需要 3 个小时以上来预处理数据集。

(vm) $ python /usr/share/torch-xla-2.0/tpu-examples/deps/dlrm/dlrm_data_pytorch.py \
    --data-generation=dataset \
    --data-set=kaggle \
    --raw-data-file=criteo-kaggle/train.txt \
    --mini-batch-size=128 \
    --memory-map \
    --test-mini-batch-size=16384 \
    --test-num-workers=4

验证预处理是否成功。

您应该会在 criteo-kaggle 目录中看到 kaggleAdDisplayChallenge_processed.npz 文件。

在预处理过的 Criteo Kaggle 数据集上运行训练脚本。

(vm) $ python /usr/share/torch-xla-2.0/tpu-examples/deps/dlrm/dlrm_tpu_runner.py \
    --arch-sparse-feature-size=16 \
    --arch-mlp-bot="13-512-256-64-16" \
    --arch-mlp-top="512-256-1" \
    --data-generation=dataset \
    --data-set=kaggle \
    --raw-data-file=criteo-kaggle/train.txt \
    --processed-data-file=criteo-kaggle/kaggleAdDisplayChallenge_processed.npz \
    --loss-function=bce \
    --round-targets=True \
    --learning-rate=0.1 \
    --mini-batch-size=128 \
    --print-freq=1024 \
    --print-time \
    --test-mini-batch-size=16384 \
    --test-num-workers=4 \
    --memory-map \
    --test-freq=101376 \
    --use-tpu \
    --num-indices-per-lookup=1 \
    --num-indices-per-lookup-fixed \
    --tpu-model-parallel-group-len 8 \
    --tpu-metrics-debug \
    --tpu-cores=8

训练应在 2 个多小时内完成，且准确率高于 78.75%。

清理

使用您创建的资源后，请进行清理，以免您的帐号产生不必要的费用：

断开与 Compute Engine 实例的连接（如果您尚未这样做）：
```
(vm) $ exit
```
您的提示符现在应为 user@projectname，表明您位于 Cloud Shell 中。
在 Cloud Shell 中，使用 Google Cloud CLI 删除 Compute Engine 实例：
```
$ gcloud compute instances delete dlrm-tutorial --zone=us-central1-a
```

使用 Google Cloud CLI 删除 Cloud TPU 资源。

$ gcloud compute tpus delete dlrm-tutorial --zone=us-central1-a

后续步骤

试用 PyTorch Colab：