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

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

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

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

费用

本教程使用 Google Cloud 的以下收费组件:

  • Compute Engine
  • Cloud TPU

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

准备工作

如果您计划在 TPU Pod 切片上训练,请务必查看此文档,该文档介绍了在 Pod 切片上训练时的特殊注意事项。

在开始学习本教程之前,请按照以下步骤检查您的 Google Cloud 项目是否已正确设置。

  1. 打开一个 Cloud Shell 窗口。

    打开 Cloud Shell

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

    export PROJECT_ID=project-id
    
  3. 配置 gcloud 命令行工具,以使用要在其中创建 Cloud TPU 的项目。

    gcloud config set project ${PROJECT_ID}
    
  4. 为 Cloud TPU 项目创建服务帐号。

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

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

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

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

    gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 -b on gs://bucket-name
    

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

  6. 使用 ctpu up 命令启动 Compute Engine 虚拟机。此示例将地区设置为 europe-west4-a,但您可以将其设置为计划用于 Compute Engine 虚拟机和 Cloud TPU 的任何地区。

    ctpu up --vm-only \
     --disk-size-gb=300 \
     --machine-type=n1-standard-8 \
     --zone=europe-west4-a \
     --tf-version=2.3.1 \
     --name=transformer-tutorial
    

    命令标志说明

    vm-only
    仅创建虚拟机。默认情况下,ctpu up 命令会同时创建虚拟机和 Cloud TPU。
    disk-size-gb
    虚拟机的磁盘大小(以 GB 为单位)。
    machine_type
    ctpu up 命令创建的虚拟机的机器类型
    zone
    拟在其中创建 Cloud TPU 的区域
    tpu-size
    要创建的 Cloud TPU 的类型
    tf-version
    在虚拟机上安装的 Tensorflow ctpu 的版本。
    name
    要创建的 Cloud TPU 的名称。

    如需详细了解 CTPU 实用程序,请参阅 CTPU 参考文档

  7. 此时会显示您指定的配置。输入 y 批准或输入 n 取消。

  8. ctpu up 命令执行完毕后,验证 shell 提示符已从 username@projectname 更改为 username@vm-name。此变化表明您现已登录 Compute Engine 虚拟机。

    gcloud compute ssh transformer-tutorial --zone=europe-west4-a
    

在您继续按照这些说明操作时,请在虚拟机会话窗口中运行以 (vm)$ 开头的每个命令。

生成训练数据集

在您的 Compute Engine 虚拟机上:

  1. 创建以下几个环境变量:将 bucket-name 替换为您的存储分区名称。

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    
    (vm)$ export GCS_DATA_DIR=$STORAGE_BUCKET/data/transformer
    (vm)$ export PARAM_SET=big
    (vm)$ export MODEL_DIR=$STORAGE_BUCKET/transformer/model_$PARAM_SET
    (vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/models"
    (vm)$ export DATA_DIR=${HOME}/transformer/data
    (vm)$ export VOCAB_FILE=${DATA_DIR}/vocab.ende.32768
    
  2. 将目录切换到训练目录:

    (vm)$ cd /usr/share/models/official/nlp/transformer
  3. 设置以下数据集环境变量:

    (vm)$ export GCS_DATA_DIR=${STORAGE_BUCKET}/data/transformer
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}
    
  4. 下载并准备数据集

    (vm)$ python3 data_download.py --data_dir=${DATA_DIR}
    (vm)$ gsutil cp -r ${DATA_DIR} ${GCS_DATA_DIR}
    

    data_download.py 下载并预处理训练和评估 WMT 数据集。数据下载并提取后,训练数据用于生成一个子词法单元词汇表。评估和训练字符串已进行词法单元处理,生成的数据会进行分片、重排并保存为 TFRecord。

    下载了 1.75GB 的压缩数据。原始文件(压缩文件、提取文件和合并文件)总计占用了 8.4 GB 的磁盘空间。生成的 TFRecord 和词汇表文件大小为 722 MB。脚本运行大约需要 40 分钟,大部分时间用于下载,大约 15 分钟用于预处理。

在单个 Cloud TPU 上训练英德翻译模型

在 Compute Engine 虚拟机上运行以下命令:

  1. 运行以下命令以创建 Cloud TPU。

    (vm)$ ctpu up --tpu-only \
      --tpu-size=v3-8  \
      --zone=europe-west4-a \
      --tf-version=2.3.1 \
      --name=transformer-tutorial
    

    命令标志说明

    tpu-only
    仅创建 Cloud TPU。默认情况下,ctpu up 命令会同时创建虚拟机和 Cloud TPU。
    tpu-size
    用于指定 Cloud TPU 的类型,例如 v3-8。
    zone
    拟在其中创建 Cloud TPU 的地区。该地区应与 Compute Engine 虚拟机所用的地区相同。例如 europe-west4-a
    tf-version
    在虚拟机上安装的 Tensorflow ctpu 的版本。
    name
    要创建的 Cloud TPU 的名称。

    如需详细了解 CTPU 实用程序,请参阅 CTPU 参考文档

    此时会显示您指定的配置。输入 y 批准或输入 n 取消。

    您会看到一条消息:Operation success; not ssh-ing to Compute Engine VM due to --tpu-only flag。由于您之前已完成 SSH 密钥传播,因此可以忽略此消息。

  2. 设置 Cloud TPU 名称变量。这可以是您在 ctpu up 中使用 --name 参数指定的名称,也可以是默认名称,即您的用户名:

    (vm)$ export TPU_NAME=transformer-tutorial
    
  3. 运行训练脚本:

    (vm)$ python3 transformer_main.py \
        --tpu=${TPU_NAME} \
        --model_dir=${MODEL_DIR} \
        --data_dir=${GCS_DATA_DIR} \
        --vocab_file=${GCS_DATA_DIR}/vocab.ende.32768 \
        --bleu_source=${GCS_DATA_DIR}/newstest2014.en \
        --bleu_ref=${GCS_DATA_DIR}/newstest2014.de \
        --batch_size=6144 \
        --train_steps=2000 \
        --static_batch=true \
        --use_ctl=true \
        --param_set=big \
        --max_length=64 \
        --decode_batch_size=32 \
        --decode_max_length=97 \
        --padded_decode=true \
        --distribution_strategy=tpu

    命令标志说明

    tpu
    Cloud TPU 的名称。这通过指定环境变量 (TPU_NAME) 进行设置。
    model_dir
    在模型训练期间存储检查点和摘要的目录。如果指定的文件夹不存在,此程序会自行创建。 使用 Cloud TPU 时,model_dir 必须是 Cloud Storage 路径 ()。您可以重复使用现有的文件夹来加载当前检查点数据和存储其他检查点,只要先前的检查点是使用相同大小的 TPU 和 TensorFlow 版本创建的即可。
    data_dir
    训练输入的 Cloud Storage 路径。在此示例中,该路径设置为 fake_imagenet 数据集。
    vocab_file
    包含要翻译的词汇的文件。
    bleu_source
    包含翻译源句的文件。
    bleu_ref
    包含翻译句子引用的文件。
    train_steps
    训练模型的步数。一个步骤处理一批数据。包括前向传导和反向传播。
    batch_size
    训练批次大小。
    static_batch
    指定数据集中的批次是否具有静态形状。
    use_ctl
    指定脚本是否通过自定义训练循环运行。
    param_set
    在创建和训练模型时使用的参数集。这些参数定义输入形状、模型配置和其他设置。
    max_length
    数据集中样本的最大长度。
    decode_batch_size
    用于在 Cloud TPU 上实现转换器自动回归解码的全局批次大小。
    decode_max_length
    解码/评估数据的最大序列长度。供转换器在 Cloud TPU 上进行自动回归解码,以最大限度地减少所需的数据填充量。
    padded_decode
    指定是否使用填充到 decode_max_length 的输入数据运行自动回归解码。Tor TPU/XLA-GPU 运行,由于静态形状要求,必须设置此标记。
    distribution_strategy
    要在 Cloud TPU 上训练 ResNet 模型,请将 distribution_strategy 设置为 tpu

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

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

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

INFO:tensorflow:Writing to file /tmp/tmpf2gn8vpa
I1125 21:22:30.734232 140328028010240 translate.py:182] Writing to file /tmp/tmpf2gn8vpa
I1125 21:22:42.785628 140328028010240 transformer_main.py:121] Bleu score (uncased): 0.01125154594774358
I1125 21:22:42.786558 140328028010240 transformer_main.py:122] Bleu score (cased): 0.01123994225054048

计算模型评估期间的 BLEU 得分

在模型评估时,使用这些标志来计算 BLEU:

  • --bleu_source:包含要翻译的文本的文件路径。
  • --bleu_ref:包含参考翻译的文件的路径。

此时,您可以结束本教程并清理 GCP 资源,也可以进一步了解如何在 Cloud TPU Pod 上运行模型。

使用 Cloud TPU Pod 扩缩模型

您可以使用 Cloud TPU Pod 扩缩模型,以便更快获得结果。这一完全受支持的 Transformer 模型可与以下 Pod 切片配合使用:

  • v2-32
  • v3-32
  1. 在您的 Cloud Shell 中,使用您在设置 Cloud TPU 时所用的 tpu-only--zone 标志运行 ctpu delete。这将仅删除您的 Cloud TPU。

     (vm)$ ctpu delete --tpu-only --zone=europe-west4-a
     

  2. 运行 ctpu up 命令,并使用 tpu-size 参数指定要使用的 Pod 切片。例如,以下命令使用 v2-32 Pod 切片。

    (vm)$ ctpu up --tpu-only \
    --tpu-size=v2-32  \
    --zone=europe-west4-a \
    --tf-version=2.3.1 \
    --name=transformer-tutorial

    命令标志说明

    tpu-only
    仅创建 Cloud TPU。默认情况下,ctpu up 命令会同时创建虚拟机和 Cloud TPU。
    tpu-size
    用于指定 Cloud TPU 的类型,例如 v3-8。
    zone
    拟在其中创建 Cloud TPU 的地区。该地区应与 Compute Engine 虚拟机所用的地区相同。例如 europe-west4-a
    tf-version
    在虚拟机上安装的 Tensorflow ctpu 的版本。
    name
    要创建的 Cloud TPU 的名称。

    如需详细了解 CTPU 实用程序,请参阅 CTPU 参考文档

    gcloud compute ssh transformer-tutorial --zone=europe-west4-a
    
  3. 导出 TPU 名称:

    (vm)$ export TPU_NAME=transformer-tutorial
    
  4. 导出模型目录变量:

    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/transformer/model_${PARAM_SET}_pod
    
  5. 将目录切换到训练目录:

    (vm)$ cd /usr/share/models/official/nlp/transformer
    
  6. 运行 Pod 训练脚本:

    (vm)$ python3 transformer_main.py \
         --tpu=${TPU_NAME} \
         --model_dir=${MODEL_DIR} \
         --data_dir=${GCS_DATA_DIR} \
         --vocab_file=${GCS_DATA_DIR}/vocab.ende.32768 \
         --bleu_source=${GCS_DATA_DIR}/newstest2014.en \
         --bleu_ref=${GCS_DATA_DIR}/newstest2014.de \
         --batch_size=24576 \
         --train_steps=2000 \
         --static_batch=true \
         --use_ctl=true \
         --param_set=big \
         --max_length=64 \
         --decode_batch_size=32 \
         --decode_max_length=97 \
         --padded_decode=true \
         --steps_between_evals=2000 \
         --distribution_strategy=tpu
    

    命令标志说明

    tpu
    Cloud TPU 的名称。这通过指定环境变量 (TPU_NAME) 进行设置。
    model_dir
    在模型训练期间存储检查点和摘要的目录。如果指定的文件夹不存在,此程序会自行创建。 使用 Cloud TPU 时,model_dir 必须是 Cloud Storage 路径 (`gs://...`)。您可以重复使用现有的文件夹来加载当前检查点数据和存储其他检查点,只要先前的检查点是使用相同大小的 TPU 和相同 Tensorflow 版本创建的即可。
    data_dir
    训练输入的 Cloud Storage 路径。在此示例中,该路径设置为 fake_imagenet 数据集。
    vocab_file
    包含要翻译的词汇的文件。
    bleu_source
    包含翻译源句的文件。
    bleu_ref
    包含翻译句子引用的文件。
    batch_size
    训练批次大小。
    train_steps
    训练模型的步数。一个步骤处理一批数据。包括前向传导和反向传播。
    static_batch
    指定数据集中的批次是否具有静态形状。
    use_ctl
    指定脚本是否通过自定义训练循环运行。
    param_set
    在创建和训练模型时使用的参数集。这些参数定义输入形状、模型配置和其他设置。
    max_length
    数据集中样本的最大长度。
    decode_batch_size
    用于在 Cloud TPU 上实现转换器自动回归解码的全局批次大小。
    decode_max_length
    解码/评估数据的最大序列长度。供转换器在 Cloud TPU 上进行自动回归解码,以最大限度地减少所需的数据填充量。
    padded_decode
    指定是否使用填充到 decode_max_length 的输入数据运行自动回归解码。Tor TPU/XLA-GPU 运行,由于静态形状要求,必须设置此标记。
    steps_between_evals
    在评估之间运行的训练步数。
    distribution_strategy
    要在 TPU 上训练 ResNet 模型,请将 distribution_strategy 设置为 tpu

此训练脚本需训练 2000 个步骤,并且每 2000 个步骤运行一次评估。在 v2-32 Cloud TPU Pod 上,这种特定的训练和评估大约需要 8 分钟。

要训练到收敛,请将 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 Platform 帐号产生费用,请执行以下操作:

1.在您的 Cloud Shell 中,使用您在设置 Cloud TPU 时所用的 --zone 标志运行 ctpu delete,以删除 Compute Engine 虚拟机和 Cloud TPU:
    $ ctpu delete --zone=europe-west4-a \
      --name=transformer-tutorial
    
  1. 运行以下命令以验证 Compute Engine 虚拟机和 Cloud TPU 均已关停:

    $ ctpu status --zone=europe-west4-a
    

    删除操作可能需要几分钟时间才能完成。 如下所示的响应表明不再有已分配的实例:

    2018/04/28 16:16:23 WARNING: Setting zone to "europe-west4-a"
    No instances currently exist.
            Compute Engine VM:     --
            Cloud TPU:             --
    
  2. 如下所示运行 gsutil,将 bucket-name 替换为您为本教程创建的 Cloud Storage 存储分区的名称:

    $ gsutil rm -r gs://bucket-name
    

后续步骤

在本教程中,您已使用示例数据集训练 Transformer 模型。此训练的结果(在大多数情况下)不能用于推断。要使用模型进行推断,您可以在公开提供的数据集或您自己的数据集上训练数据。在 Cloud TPU 上训练的模型要求数据集采用 TFRecord 格式。

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

超参数调节

如需使用数据集提升模型的性能,您可以调节模型的超参数。您可以在 GitHub 上寻找所有 TPU 支持模型通用的超参数的相关信息。您可以在每个模型的源代码中寻找模型专用超参数的相关信息。如需详细了解超参数调节,请参阅超参数调节概览使用超参数调节服务调节超参数

推理

训练模型后,您可以使用该模型进行推断(也称为预测)。AI Platform 是一款基于云的解决方案,用于开发、训练部署机器学习模型。部署模型后,您可以使用 AI Platform Prediction 服务