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

本教程介绍如何训练 DLRM 和 DCN v2 排名模型，用于预测点击率 (CTR) 等任务。查看以下语言版本的备注设置以运行 DLRM 或 DCN 模型，了解如何设置参数来训练 DLRM 或 DCN v2 排名模型。

模型输入是数值特征和分类特征，输出是标量（例如点击概率）。模型可以 Cloud TPU。深度排名模型都属于内存密集型模型（对于嵌入，表和查询），以及深度网络 (MLP) 的计算密集型应用。TPU 两者兼顾

该模型将 TPUEmbedding 层用于分类特征。TPU 嵌入支持具有快速查找功能的大型嵌入表，会根据 TPU Pod 的大小线性扩缩。您最多只能使用 90 GB 的嵌入表用于 TPU v3-8，5.6 TB 用于 v3-512 Pod，22.4 TB 用于 v3-2048 TPU Pod。

模型代码位于 TensorFlow Recommenders 库中，其中介绍了输入流水线、配置和训练循环， TensorFlow Model Garden。

目标

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

费用

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

Compute Engine
Cloud TPU
Cloud Storage

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

准备工作

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

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

Go to project selector
确保您的 Google Cloud 项目已启用结算功能。

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

设置资源

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

打开一个 Cloud Shell 窗口。

打开 Cloud Shell
为项目 ID 创建一个变量。

    export PROJECT_ID=project-id

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

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

gcloud config set project ${PROJECT_ID}

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

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

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

该命令会返回一个 Cloud TPU 服务账号，其格式如下：

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

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

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

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

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

使用 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 的[区域](https://cloud.google.com/tpu/docs/types-zones?hl=zh-cn)。

`accelerator-type`

加速器类型指定要创建的 Cloud TPU 的版本和大小。 如需详细了解每个 TPU 版本支持的加速器类型，请参阅 [TPU 版本](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm?hl=zh-cn#versions)。

`version`

Cloud TPU [软件版本](https://cloud.google.com/tpu/docs/supported-tpu-versions?hl=zh-cn#tpu_software_versions)。

使用 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 TB 和 Criteo Kaggle。本教程通过设置标志 use_synthetic_data=True 使用合成数据进行训练。

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

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

安装必需的软件包。

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

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

运行训练脚本。它使用类似 Criteo 的虚构数据集来训练 DLRM 模型。训练大约需要 20 分钟。

auto
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 分钟。完成后，您将看到如下所示的消息：

在这里插入图片描述

清除数据

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

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

(vm)$ exit

你的提示现在应为username@projectname，表明你处于运行命令的 Cloud Shell 实际效果

删除您的 Cloud TPU 资源。

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

通过运行 gcloud compute tpus execution-groups list 验证资源是否已删除。删除操作可能需要几分钟时间才能完成。以下命令的输出不应包含本教程中创建的任何资源：

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

使用 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。如需了解特定于模型的超参数，请参阅源代码每个代码模型。如需详细了解超参数调优，请参阅概览超参数调优和 Tune 超参数。

推断

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

文章来源：google cloud

推荐阅读

在 Cloud TPU 上训练 Mask RCNN (TF 2.x)
在 Cloud TPU 上训练 ShapeMask (TF 2.x)
在 Cloud TPU 上训练 RetinaNet (TF 2.x)

更多芯擎AI开发板干货请关注芯擎AI开发板专栏。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。