在 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. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  5. Make sure that billing is enabled for your Google Cloud project.

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

设置资源

本部分介绍如何设置本教程使用的 Cloud Storage 存储桶、虚拟机和 Cloud TPU 资源。

  1. 打开一个 Cloud Shell 窗口。

    打开 Cloud Shell

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

    export PROJECT_ID=project-id
  3. 配置 Google Cloud CLI 以使用要在其中创建项目的项目 Cloud TPU。

    如需详细了解 gcloud 命令,请参阅 Google Cloud CLI 参考文档

    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 存储桶,其中 --location 选项指定应创建存储桶的区域。如需详细了解可用区和区域,请参阅类型和可用区

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

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

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

  6. 使用 gcloud 命令启动 Compute Engine 虚拟机和 Cloud TPU。

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

    命令标志说明

    zone
    拟在其中创建 Cloud TPU 的区域
    accelerator-type
    加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。 如需详细了解每个 TPU 版本支持的加速器类型,请参阅 TPU 版本
    version
    Cloud TPU 软件版本
  7. 使用 SSH 连接到 Compute Engine 实例。连接到网络后 您的 Shell 提示符会从 username@projectname 更改为 username@vm-name:

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

设置 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 名称设置环境变量。

  (vm)$ export TPU_NAME=local
  

训练应用预期能够访问您在 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. 切换到脚本目录。

    (vm)$ cd /usr/share/tpu/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 资源。

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

    $ gcloud compute tpus tpu-vm list --zone=europe-west4-a
  4. 使用 gcloud CLI 删除 Cloud Storage 存储桶。将 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。 您可以在每个模型的源代码中寻找模型专用超参数的相关信息。如需详细了解超参数调节,请参阅超参数调节概览调节超参数

推断

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