在 Cloud TPU 上训练 Transformer (TF 2.x)

---

如果您不熟悉 Cloud TPU，请参阅 TPU 使用入门指南。

本教程介绍如何在 Cloud TPU 上训练 Transformer 模型。Transformer 是一种神经网络架构，它使用注意力机制来解决序列到序列问题。与传统的神经 seq2seq 模型不同，Transformer 不涉及循环连接。注意力机制分两个序列获知词法单元之间的依赖关系。由于注意力权重应用于序列中的所有词法单元，因此 Transformer 模型能够轻松捕获长距离依赖关系。

Transformer 的整体结构遵循标准的编码器 - 解码器模式。编码器使用自注意力来计算输入序列的表示形式。解码器将编码器输出和先前解码器输出的词法单元作为输入，按一次一个词法单元的模式生成输出序列。

该模型还会对输入和输出词法单元应用嵌入，并添加常量位置编码。位置编码会添加有关每个词法单元位置的信息。

警告：本教程使用第三方数据集。Google 不对此数据集的有效性或任何其他方面提供任何代理、担保或其他保证。

目标

• 创建 Cloud Storage 存储分区以保存数据集和模型输出。
• 下载并预处理用于训练模型的数据集。
• 运行训练作业。
• 验证输出结果。

费用

在本文档中，您将使用 Google Cloud 的以下收费组件：

• Compute Engine
• Cloud TPU

您可使用价格计算器根据您的预计使用情况来估算费用。 Google Cloud 新用户可能有资格申请免费试用。

准备工作

在开始学习本教程之前，请检查您的 Google Cloud 项目是否已正确设置。

1. 登录您的 Google Cloud 账号。如果您是 Google Cloud 新手，请创建一个账号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金，用于运行、测试和部署工作负载。

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. 确保您的 Google Cloud 项目已启用结算功能。

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. 确保您的 Google Cloud 项目已启用结算功能。

本演示使用 Google Cloud 的收费组件。请查看 Cloud TPU 价格页面估算您的费用。请务必在使用完您创建的资源以后清理这些资源，以免产生不必要的费用。

使用单个 Cloud TPU 设备进行训练

注意：您是否要使用 Cloud TPU Pod 训练模型？请参阅 TPU Pod 训练。

本部分介绍如何为单一设备训练设置 Cloud Storage 存储桶和 TPU 虚拟机。

重要提示：在同一区域和可用区中设置所有资源（Cloud TPU 虚拟机和 Cloud Storage 存储桶），以减少网络延迟和网络费用。TPU 虚拟机位于特定可用区，即区域内的细分。

1. 打开一个 Cloud Shell 窗口。

   打开 Cloud Shell

2. 为您的项目 ID 创建环境变量。

   export PROJECT_ID=project-id

3. 配置 Google Cloud CLI 以使用您要创建 Cloud TPU 的 Google Cloud 项目。

   gcloud config set project ${PROJECT_ID}

   当您第一次在新的 Cloud Shell 虚拟机中运行此命令时，系统会显示 Authorize Cloud Shell 页面。点击页面底部的 Authorize 以允许 gcloud 使用您的凭据进行 Google Cloud API 调用。

4. 为 Cloud TPU 项目创建服务账号。

   通过服务账号，Cloud TPU 服务可以访问其他 Google Cloud 服务。

   $ gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID

   该命令将返回以下格式的 Cloud TPU 服务账号：

   [e-mail]
   

5. 使用以下命令创建 Cloud Storage 存储分区：

   $ gcloud storage buckets create gs://bucket-name --project=${PROJECT_ID} --location=us-central2

   此 Cloud Storage 存储分区存储您用于训练模型的数据和训练结果。本教程中使用的 gcloud 命令会为您设置 TPU，还会为在上一步中设置的 Cloud TPU 服务账号设置默认权限。如果您需要更精细的权限，请查看访问级层权限。

在单个 Cloud TPU 上训练 Transformer 模型

1. 使用 gcloud 命令启动 Cloud TPU 虚拟机。

     $ gcloud compute tpus tpu-vm create transformer-tutorial \
        --zone=us-central2-b \
        --accelerator-type=v4-8 \
        --version=tpu-vm-tf-2.17.0-pjrt

   命令标志说明

   zone

   拟在其中创建 Cloud TPU 的区域。

   accelerator-type

   加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。 如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。

   version

   Cloud TPU 软件版本。

   注意 ：如果您有多个 Google Cloud 项目，则必须使用 --project 标志，用于指定 Google Cloud 的 ID 在其中创建 Cloud TPU 的位置。
   注意 ：您首次对项目运行 gcloud compute tpus 时 执行启动任务（例如 SSH 密钥）大约需要 5 分钟 传播和 API 启动。

2. 使用 SSH 连接到 TPU 虚拟机。连接到虚拟机后，shell 提示符会从 username@projectname 更改为 username@vm-name：

   gcloud compute tpus tpu-vm ssh transformer-tutorial --zone=us-central2-b

   要点 ：从现在开始，前缀 (vm) $ 表示您应该运行 对 Compute Engine 虚拟机实例执行操作。

3. 导出环境变量。

     (vm)$ export STORAGE_BUCKET=gs://bucket-name
     (vm)$ export SENTENCEPIECE_MODEL=sentencepiece
     (vm)$ export SENTENCEPIECE_MODEL_PATH=${STORAGE_BUCKET}/${SENTENCEPIECE_MODEL}.model
     (vm)$ export TFDS_DIR=${STORAGE_BUCKET}/tfds
     (vm)$ export PARAM_SET=big
     (vm)$ export TPU_NAME=local
     (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}
     (vm)$  export PYTHONPATH="/usr/share/tpu/models:$PYTHONPATH"

4. 创建 TPU 时，如果您将 --version 参数设置为以 -pjrt 时，请设置以下环境变量以启用 PJRT 运行时：

     (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
     (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

5. 安装 TensorFlow 要求。

   (vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt

6. 下载并预处理数据集

    (vm)$ python3 -c "import tensorflow_datasets as tfds; tfds.load('wmt14_translate/de-en', split='train+validation', shuffle_files=True, download=True)"
     (vm)$ python3 /usr/share/tpu/models/official/nlp/data/train_sentencepiece.py --output_model_path=${SENTENCEPIECE_MODEL}

7. 将数据集复制到 Cloud Storage 存储桶

     (vm)$ gcloud storage cp ${SENTENCEPIECE_MODEL}.model ${STORAGE_BUCKET}
     (vm)$ gcloud storage cp ${SENTENCEPIECE_MODEL}.vocab ${STORAGE_BUCKET}
     (vm)$ gcloud storage cp tensorflow_datasets/wmt14_translate ${TFDS_DIR}/wmt14_translate --recursive

8. 导航到训练目录

   (vm)$ cd /usr/share/tpu/models/

9. 运行训练脚本

   (vm)$ python3 official/nlp/train.py \
     --tpu=${TPU_NAME} \
     --experiment=wmt_transformer/large \
     --mode=train_and_eval \
     --model_dir=${MODEL_DIR} \
     --params_override="runtime.distribution_strategy=tpu, task.train_data.tfds_data_dir=${TFDS_DIR}, task.validation_data.tfds_data_dir=${TFDS_DIR}, task.sentencepiece_model_path=${SENTENCEPIECE_MODEL_PATH}, trainer.train_steps=10000, trainer.validation_interval=10000"

   命令标志说明

   tpu

   Cloud TPU 的名称。您可以通过指定 环境变量 (TPU_NAME)。

   experiment

   要训练的模型。

   mode

   脚本的运行模式。

   model_dir

   在训练期间存储检查点和摘要的 Cloud Storage 存储分区。您可以使用现有文件夹加载根据相同大小和 TensorFlow 版本的先前生成 TPU 创建的检查点。

   params_override

   设置模型参数。

   默认情况下，模型每完成 10,000 步后会进行一次评估。您可以 增加训练步数或指定运行评估的频率 方法是设置以下参数：

   • train.train_steps：要运行的总训练步数。
   • trainer.validation_interval：在评估之间运行的训练步数。

   在 v4-8 上训练和评估大约需要 20 分钟 Cloud TPU。训练和评估完成后，系统将显示类似如下内容的消息：

   I0208 20:57:19.309512 140039467895872 controller.py:310]  eval | step:  10000 | eval time:   69.2 sec | output:
   {'bleu_score': 19.204771518707275,
    'sacrebleu_score': 18.307039308307356,
    'validation_loss': 2.0654342}
    eval | step:  10000 | eval time:   69.2 sec | output:
   {'bleu_score': 19.204771518707275,
    'sacrebleu_score': 18.307039308307356,
    'validation_loss': 2.0654342}
   

   您现已完成单设备训练。请按照以下步骤删除单设备 TPU 资源。

10. 与 Compute Engine 实例断开连接：

    (vm)$ exit

    您的提示符现在应为 username@projectname，表明您位于 Cloud Shell 中。

11. 删除 TPU 资源。

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

    命令标志说明

    zone

    您的 Cloud TPU 所在的可用区。

此时，您可以结束本教程并清理，也可以继续并探索在 Cloud TPU Pod 上运行模型。

使用 Cloud TPU Pod 扩缩模型

在 Cloud TPU Pod 上训练模型可能需要对模型进行一些更改， 训练脚本。如需了解相关信息，请参阅在 TPU Pod 上训练。

TPU Pod 训练

重要提示：如果您已为单设备训练设置 Cloud TPU 项目、服务账号、存储桶和数据集，则可以跳到启动 TPU 资源。

重要提示：在同一区域和可用区设置 Cloud TPU 资源和 Cloud Storage 存储桶，以减少网络延迟和网络费用。Cloud TPU 位于特定可用区内，这些可用区是区域内的细分。

1. 打开一个 Cloud Shell 窗口。

   打开 Cloud Shell

2. 为项目 ID 创建一个变量。

   export PROJECT_ID=project-id

3. 配置 Google Cloud CLI 以使用要在其中创建 Cloud TPU 的项目。

   gcloud config set project ${PROJECT_ID}

   当您第一次在新的 Cloud Shell 虚拟机中运行此命令时，系统会显示 Authorize Cloud Shell 页面。点击页面底部的 Authorize 以允许 gcloud 使用您的凭据进行 API 调用。

4. 为 Cloud TPU 项目创建服务账号。

   gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID

   该命令将返回以下格式的 Cloud TPU 服务账号：

   [e-mail]
   

5. 使用以下命令创建 Cloud Storage 存储分区，或使用您之前为项目创建的存储分区：

   gcloud storage buckets create gs://bucket-name --project=${PROJECT_ID} --location=us-central1

   此 Cloud Storage 存储分区存储您用于训练模型的数据和训练结果。本教程中使用的 gcloud 命令会设置 您在 Cloud TPU 和 Cloud TPU 中 上一步。如需详细了解更精细的权限，请参阅 访问权限级别权限。

   重要提示 ：如果您创建了新的 Cloud Storage 存储桶，则需要 下载并预处理数据集，如本文档前面所述。

启动 TPU 虚拟机资源

1. 使用 gcloud 命令启动 TPU 虚拟机 Pod。这个 本教程指定了 v4-32 Pod如需了解其他 Pod 选项，请参阅 TPU 类型可用的 TPU 类型页面。

   注意：如果当前没有足够的容量来创建 TPU Pod，您可以 使用排队的资源将请求加入队列加入队列的资源让您可以在资源用尽后 可用。如需以已排队的资源形式请求 Cloud TPU 资源，请改用 gcloud compute tpus queued-resources create 命令。如需了解详情，请参阅管理队列资源。

   $ gcloud compute tpus tpu-vm create transformer-tutorial \
     --zone=us-central2-b \
     --accelerator-type=v3-32 \
     --version=tpu-vm-tf-2.17.0-pod-pjrt

   命令标志说明

   zone

   拟在其中创建 Cloud TPU 的区域。

   accelerator-type

   加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。 如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。

   version

   Cloud TPU 软件版本。

2. 使用 SSH 连接到 Compute Engine 实例。连接到网络后 您的 Shell 提示符会从 username@projectname 更改为 username@vm-name:

   gcloud compute tpus tpu-vm ssh transformer-tutorial --zone=us-central2-b

   要点：从现在起，前缀 (vm) $ 表示您应在 Compute Engine 虚拟机实例上运行该命令。

3. 安装 TensorFlow 要求。

   (vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt

设置并启动 Pod 训练

1. 导出 Cloud TPU 设置变量：

   (vm)$ export PYTHONPATH="/usr/share/tpu/models:$PYTHONPATH"
   (vm)$ export STORAGE_BUCKET=gs://bucket-name
   (vm)$ export SENTENCEPIECE_MODEL=sentencepiece
   (vm)$ export SENTENCEPIECE_MODEL_PATH=${STORAGE_BUCKET}/${SENTENCEPIECE_MODEL}.model
   (vm)$ export TFDS_DIR=${STORAGE_BUCKET}/tfds
   (vm)$ export TPU_NAME=transformer-tutorial
   (vm)$ export PARAM_SET=big
   (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}
   (vm)$ export TPU_LOAD_LIBRARY=0

2. 下载数据集

   注意 ：如果您已下载数据集，请跳过此步骤。

   (vm)$ python3 -c "import tensorflow_datasets as tfds; tfds.load('wmt14_translate/de-en', split='train+validation', shuffle_files=True, download=True)"
   (vm)$ python3 /usr/share/tpu/models/official/nlp/data/train_sentencepiece.py --output_model_path=${SENTENCEPIECE_MODEL}

3. 将数据集复制到 Cloud Storage 存储桶

   注意 ：如果您已经下载数据集并将其复制到 Cloud Storage 存储桶。

   (vm)$ gcloud storage cp ${SENTENCEPIECE_MODEL}.model ${STORAGE_BUCKET}
   (vm)$ gcloud storage cp ${SENTENCEPIECE_MODEL}.vocab ${STORAGE_BUCKET}
   (vm)$ gcloud storage cp tensorflow_datasets/wmt14_translate ${TFDS_DIR}/wmt14_translate --recursive

4. 切换到训练目录：

   (vm)$ cd /usr/share/tpu/models/

5. 运行训练脚本：

   (vm)$ python3 official/nlp/train.py \
   --tpu=${TPU_NAME} \
   --experiment=wmt_transformer/large \
   --mode=train_and_eval \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu, task.train_data.tfds_data_dir=${TFDS_DIR}, task.validation_data.tfds_data_dir=${TFDS_DIR}, task.sentencepiece_model_path=${SENTENCEPIECE_MODEL_PATH}, trainer.train_steps=10000, trainer.validation_interval=10000"

   命令标志说明

   tpu

   Cloud TPU 的名称。您可以通过指定 环境变量 (TPU_NAME)。

   experiment

   要训练的模型。

   mode

   脚本的运行模式。

   model_dir

   在训练期间存储检查点和摘要的 Cloud Storage 存储分区。您可以使用现有文件夹加载根据相同大小和 TensorFlow 版本的先前生成 TPU 创建的检查点。

   params_override

   设置模型参数。

默认情况下，该模型每 10,000 步进行一次评估。为了训练 要收敛，请将 train_steps 更改为 200000。您可以设置以下参数来增加训练步骤数，或指定评估的运行频率：

• trainer.train_steps：用于设置要运行的总训练步数。
• trainer.validation_interval：设置要运行的训练步数 评估。

在 v4-32 上训练和评估大约需要 14 分钟 Cloud TPU。训练和评估完成后，系统将显示类似如下内容的消息：

I0209 22:19:49.143219 139751309618240 controller.py:310]  eval | step:  10000 | eval time:   73.6 sec | output:
    {'bleu_score': 19.401752948760986,
     'sacrebleu_score': 18.442741330886378,
     'validation_loss': 2.0558002}
 eval | step:  10000 | eval time:   73.6 sec | output:
    {'bleu_score': 19.401752948760986,
     'sacrebleu_score': 18.442741330886378,
     'validation_loss': 2.0558002}

重要提示：完成训练和评估后，请务必清理，这样就不需要为额外的 TPU 使用付费。

此训练脚本需训练 20000 个步骤，并且每 2000 个步骤运行一次评估。在 v3-32 Cloud TPU Pod 上，这种特定的训练和评估大约需要 8 分钟。训练和评估完成后，系统将显示类似如下内容的消息：

INFO:tensorflow:Writing to file /tmp/tmpdmlanxcf
I0218 21:09:19.100718 140509661046592 translate.py:184] Writing to file /tmp/tmpdmlanxcf
I0218 21:09:28.043537 140509661046592 transformer_main.py:118] Bleu score (uncased): 1.799112930893898
I0218 21:09:28.043911 140509661046592 transformer_main.py:119] Bleu score (cased): 1.730366237461567

要训练到收敛，请将 train_steps 更改为 200000。您可以设置以下参数来增加训练步骤数，或指定评估的运行频率：

• --train_steps：用于设置要运行的总训练步数。
• --steps_between_evals：在评估之间运行的训练步数。

训练和评估完成后，系统将显示类似如下内容的消息：

0509 00:27:59.984464 140553148962624 translate.py:184] Writing to file /tmp/tmp_rk3m8jp
I0509 00:28:11.189308 140553148962624 transformer_main.py:119] Bleu score (uncased): 1.3239131309092045
I0509 00:28:11.189623 140553148962624 transformer_main.py:120] Bleu score (cased): 1.2855342589318752

清除数据

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的项目，或者保留项目但删除各个资源。

1. 断开与 Compute Engine 实例的连接（如果您尚未这样做）：

   (vm)$ exit

   您的提示符现在应为 username@projectname，表明您位于 Cloud Shell 中。

2. 删除您的 Cloud TPU 和 Compute Engine 资源。

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

3. 如下所示，运行 gcloud CLI，将 bucket-name 替换为 您为本教程创建的 Cloud Storage 存储桶的名称：

   警告：删除存储分区后，所有训练数据都将丢失，因此，请仅在完成运行本教程后才执行此步骤。

   $ gcloud storage rm gs://bucket-name --recursive

后续步骤

TensorFlow Cloud TPU 教程通常使用示例数据集来训练模型。此训练的结果不能用于推断。如需使用模型进行推理，您可以在公开提供的数据集或您自己的数据集上训练数据。在 Cloud TPU 上训练的 TensorFlow 模型 通常需要将数据集 TFRecord 格式。

您可以使用数据集转换工具 示例，用于转换图片 转换为 TFRecord 格式。如果您未使用图片分类模型，则必须自行将数据集转换为 TFRecord 格式。如需了解详情，请参阅 TFRecord 和 tf.Example。

超参数调节

如需使用数据集提升模型的性能，您可以调节模型的超参数。您可以找到所有用户共有的超参数的相关信息 支持 TPU 的 GitHub。 如需了解特定于模型的超参数，请参阅源代码 每个代码 模型。如需详细了解超参数调优，请参阅概览 超参数调优 和 Tune 超参数。

推断

训练模型后，您可以使用该模型进行推理（也称为预测）。您可以使用 Cloud TPU 推断转换器 工具来准备和优化 在 Cloud TPU v5e 上用于推理的 TensorFlow 模型。如需详细了解 Cloud TPU v5e 上的推理，请参阅 Cloud TPU v5e 推理简介。