在 Cloud TPU 上训练 DLRM 和 DCN (TF 2.x)


本教程介绍如何训练 DLRM 和 DCN v2 排名模型,该模型可用于点击率 (CTR) 预测等任务。请参阅进行设置以运行 DLRM 或 DCN 模型中的说明,了解如何设置参数以训练 DLRM 或 DCN v2 排名模型。

模型输入是数值特征和分类特征,而输出是标量(例如点击概率)。您可以在 Cloud TPU 上训练和评估该模型。对深度网络 (MLP) 而言,深度排名模型具有内存密集(嵌入表/查询)和计算密集的特征。TPU 是专门针对这两个特征设计的。

该模型将 TPUembedding 层用于分类特征。TPU 嵌入支持可快速查询的大型嵌入表,嵌入表的大小随 TPU pod 的大小线性扩缩。TPU v3-8 可使用多达 90 GB 的嵌入表,v3-512 Pod 可使用 5.6 TB,v3-2048 TPU Pod 可使用 22.4 TB。

模型代码位于 TensorFlow Recommenders 库中,而输入流水线、配置和训练循环位于 TensorFlow Model Garden 中。

目标

  • 设置训练环境
  • 使用合成数据运行训练作业
  • 验证输出结果

费用

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

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

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

准备工作

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

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

    转到“项目选择器”

  3. 确保您的 Google Cloud 项目已启用结算功能

  4. 在 Google Cloud Console 中的项目选择器页面上,选择或创建一个 Google Cloud 项目

    转到“项目选择器”

  5. 确保您的 Google Cloud 项目已启用结算功能

  6. 本演示使用 Google Cloud 的收费组件。请查看 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 服务账号:

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

  5. 使用以下命令创建 Cloud Storage 存储桶,其中 -l 选项指定应创建存储桶的区域。如需详细了解可用区和区域,请参阅类型和可用区

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

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

    存储桶位置必须要与 Compute Engine(虚拟机)和 Cloud TPU 节点位于同一区域。

  6. 使用 gcloud 命令启动 Compute Engine 虚拟机和 Cloud TPU。 使用的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需详细了解这两种虚拟机架构,请参阅系统架构。 如需详细了解 gcloud 命令,请参阅 gcloud 参考文档

    TPU 虚拟机

    $ gcloud compute tpus tpu-vm create dlrm-dcn-tutorial \
    --zone=europe-west4-a \
    --accelerator-type=v3-8 \
    --version=tpu-vm-tf-2.16.1-se

    命令标志说明

    zone
    您计划在其中创建 Cloud TPU 的区域
    accelerator-type
    加速器类型指定要创建的 Cloud TPU 的版本和大小。如需详细了解每个 TPU 版本支持的加速器类型,请参阅 TPU 版本
    version
    Cloud TPU 软件版本

    TPU 节点

    $ gcloud compute tpus execution-groups create \
    --name=dlrm-dcn-tutorial \
    --zone=europe-west4-a \
    --disk-size=300 \
    --machine-type=n1-standard-8 \
    --tf-version=2.12.0
    

    命令标志说明

    vm-only
    仅创建虚拟机。默认情况下,gcloud compute tpus execution-groups 命令会同时创建虚拟机和 Cloud TPU。
    name
    要创建的 Cloud TPU 的名称。
    zone
    您计划在其中创建 Cloud TPU 的区域
    disk-size
    通过 gcloud compute tpus execution-groups 命令创建的虚拟机的硬盘大小(以 GB 为单位)。
    machine-type
    要创建的 Compute Engine 虚拟机的机器类型
    tf-version
    TensorFlow ctpu 的版本会安装在虚拟机上。
  7. 如果您未自动登录 Compute Engine 实例,请通过运行以下 ssh 命令进行登录。登录虚拟机后,shell 提示符会从 username@projectname 更改为 username@vm-name

    TPU 虚拟机

    gcloud compute tpus tpu-vm ssh dlrm-dcn-tutorial --zone=europe-west4-a
    

    TPU 节点

    gcloud compute ssh dlrm-dcn-tutorial --zone=europe-west4-a
    

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

设置 Cloud Storage 存储桶变量

设置以下环境变量,将 bucket-name 替换为 Cloud Storage 存储桶的名称:

(vm)$ export STORAGE_BUCKET=gs://bucket-name
(vm)$ export PYTHONPATH="/usr/share/tpu/models/:${PYTHONPATH}"
(vm)$ export EXPERIMENT_NAME=dlrm-exp

为 TPU 名称设置环境变量。

TPU 虚拟机

(vm)$ export TPU_NAME=local

TPU 节点

(vm)$ export TPU_NAME=dlrm-dcn-tutorial

训练应用预期能够访问您在 Cloud Storage 中的训练数据。在训练期间,训练应用还会使用您的 Cloud Storage 存储桶来存储检查点。

进行设置以使用合成数据运行 DLRM 或 DCN 模型

该模型可以使用各种数据集进行训练。最常用的两个数据集是 Criteo TBCriteo Kaggle。本教程通过设置标志 use_synthetic_data=True 使用合成数据进行训练。

合成数据集仅用于了解如何使用 Cloud TPU 和验证端到端性能。准确率数字和保存的模型没有意义。

请访问 Criteo TerabyteCriteo Kaggle 网站,了解如何下载和预处理这些数据集。

  1. 安装必需的软件包。

    (vm)$ pip3 install tensorflow-recommenders
    (vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt
    
  2. 切换到脚本目录。

    TPU 虚拟机

    (vm)$ cd /usr/share/tpu/models/official/recommendation/ranking

    TPU 节点

    (vm)$ cd /usr/share/models/official/recommendation/ranking
  3. 运行训练脚本。此脚本使用类似于 Criteo 的虚构数据集来训练 DLRM 模型。训练大约需要 20 分钟。

    export EMBEDDING_DIM=32
    
    python3 train.py --mode=train_and_eval \
         --model_dir=${STORAGE_BUCKET}/model_dirs/${EXPERIMENT_NAME} --params_override="
         runtime:
             distribution_strategy: 'tpu'
         task:
             use_synthetic_data: true
             train_data:
                 input_path: '${DATA_DIR}/train/*'
                 global_batch_size: 16384
             validation_data:
                 input_path: '${DATA_DIR}/eval/*'
                 global_batch_size: 16384
             model:
                 num_dense_features: 13
                 bottom_mlp: [512,256,${EMBEDDING_DIM}]
                 embedding_dim: ${EMBEDDING_DIM}
                 top_mlp: [1024,1024,512,256,1]
                 interaction: 'dot'
                 vocab_sizes: [39884406, 39043, 17289, 7420, 20263, 3, 7120, 1543, 63,
                     38532951, 2953546, 403346, 10, 2208, 11938, 155, 4, 976, 14,
                     39979771, 25641295, 39664984, 585935, 12972, 108, 36]
         trainer:
             use_orbit: false
             validation_interval: 1000
             checkpoint_interval: 1000
             validation_steps: 500
             train_steps: 1000
             steps_per_loop: 1000
         "
    

此训练在 v3-8 TPU 上运行大约 10 分钟。完成后,您将看到如下所示的消息:

I0621 21:32:58.519792 139675269142336 tpu_embedding_v2_utils.py:907] Done with log of TPUEmbeddingConfiguration.
I0621 21:32:58.540874 139675269142336 tpu_embedding_v2.py:389] Done initializing TPU Embedding engine.
1000/1000 [==============================] - 335s 335ms/step - auc: 0.7360 - accuracy: 0.6709 - prediction_mean: 0.4984
- label_mean: 0.4976 - loss: 0.0734 - regularization_loss: 0.0000e+00 - total_loss: 0.0734 - val_auc: 0.7403
- val_accuracy: 0.6745 - val_prediction_mean: 0.5065 - val_label_mean: 0.4976 - val_loss: 0.0749
- val_regularization_loss: 0.0000e+00 - val_total_loss: 0.0749

Model: "ranking"
_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
tpu_embedding (TPUEmbedding) multiple                  1
_________________________________________________________________
mlp (MLP)                    multiple                  154944
_________________________________________________________________
mlp_1 (MLP)                  multiple                  2131969
_________________________________________________________________
dot_interaction (DotInteract multiple                  0
_________________________________________________________________
ranking_1 (Ranking)          multiple                  0
=================================================================
Total params: 2,286,914
Trainable params: 2,286,914
Non-trainable params: 0
_________________________________________________________________
I0621 21:43:54.977140 139675269142336 train.py:177] Train history: {'auc': [0.7359596490859985],
'accuracy': [0.67094486951828], 'prediction_mean': [0.4983849823474884], 'label_mean': [0.4975697994232178],
'loss': [0.07338511198759079], 'regularization_loss': [0], 'total_loss': [0.07338511198759079],
'val_auc': [0.7402724623680115], 'val_accuracy': [0.6744520664215088], 'val_prediction_mean': [0.5064718723297119],
'val_label_mean': [0.4975748658180237], 'val_loss': [0.07486172765493393],
'val_regularization_loss': [0], 'val_total_loss': [0.07486172765493393]}

清理

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

  1. 断开与 Compute Engine 实例的连接(如果您尚未这样做):

    (vm)$ exit
    

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

  2. 删除您的 Cloud TPU 和 Compute Engine 资源。 用于删除资源的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需了解详情,请参阅系统架构

    TPU 虚拟机

    $ gcloud compute tpus tpu-vm delete dlrm-dcn-tutorial \
    --zone=europe-west4-a
    

    TPU 节点

    $ gcloud compute tpus execution-groups delete dlrm-dcn-tutorial \
    --zone=europe-west4-a
    
  3. 通过运行 gcloud compute tpus execution-groups list 验证资源是否已删除。删除操作可能需要几分钟时间才能完成。以下命令的输出不应包含本教程中创建的任何资源:

    TPU 虚拟机

    $ gcloud compute tpus tpu-vm list --zone=europe-west4-a

    TPU 节点

    $ gcloud compute tpus execution-groups list --zone=europe-west4-a
  4. 使用 gsutil 删除 Cloud Storage 存储桶。将 bucket-name 替换为您的 Cloud Storage 存储桶的名称。

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

后续步骤

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

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

超参数调优

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

推理

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