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

概览

这是神经协同过滤论文中所述的采用神经矩阵分解 (NeuMF) 模型的神经协同过滤 (NCF) 框架的实现。目前的实现基于作者的 NCF 代码中的代码以及 MLPerf 代码库中的 Stanford 实现。

NCF 是一个用于推荐的协同过滤的通用框架，在该框架中，神经网络架构用于为用户-推荐项交互建模。与传统模型不同，NCF 并不针对用户和推荐项潜在特征的内积使用矩阵分解 (MF)。它将内积替换为多层感知机，该感知机可以从数据中学习任意函数。

NCF 有两种实现方式：泛化矩阵分解 (GMF) 和多层感知机 (MLP)。GMF 应用线性内核来为潜在特征交互建模，MLP 使用非线性内核从数据中学习交互函数。NeuMF 是 GMF 和 MLP 的融合模型，可以更好地为复杂的用户-推荐项交互建模，并统一了 MLP 的线性和 MLP 的非线性在为用户-推荐项潜在结构建模方面的优势。NeuMF 允许 GMF 和 MLP 学习不同的嵌入，并通过连接最后一个隐藏层将两个模型组合在一起。neumf_model.py 定义了架构细节。

以下说明假设您已熟悉如何在 Cloud TPU 上训练模型。如果您刚开始接触 Cloud TPU，请参阅快速入门，了解基本介绍。

数据集

MovieLens 数据集用于模型训练和评估。具体来说，我们使用两个数据集：ml-1m（MovieLens 一百万的简称）和 ml-20m（MovieLens 两千万的简称）。

ml-1m

ml-1m 数据集包含 2000 年加入 MovieLens 的 6040 名用户对约 3706 部电影的 1000209 个匿名评分。所有评分都包含在文件“ratings.dat”中，没有标题行，并且采用以下格式：

UserID::MovieID::Rating::Timestamp

UserID 的范围介于 1 到 6040 之间。
MovieID 的范围介于 1 到 3952 之间。
评分采用 5 星制（仅限全星评分）。

ml-20m

ml-20m 数据集包含 138493 名用户对 26744 部电影的 20000263 个评分。所有评分都包含在文件“ratings.csv”中。在此文件中，标题行之后的每一行代表一个用户对一部电影的一个评分，并采用以下格式：

userId,movieId,rating,timestamp

此文件中的行先按 userId 排序，在同一个用户中，再按 movieId 排序。评分采用 5 星制，按 0.5 星递增（0.5 星 - 5.0 星）。两个数据集中的时间戳均以从世界协调时间 (UTC) 1970 年 1 月 1 日零点起计算的秒数表示。每个用户至少有 20 个评分。

目标

创建 Cloud Storage 存储桶以保存数据集和模型输出
准备 MovieLens 数据集
设置 Compute Engine 虚拟机和 Cloud TPU 节点以进行训练和评估
运行训练和评估

费用

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

Compute Engine
Cloud TPU
Cloud Storage

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

准备工作

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

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

在 Google Cloud Console 中的项目选择器页面上，选择或创建一个 Google Cloud 项目。

转到“项目选择器”

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

在 Google Cloud Console 中的项目选择器页面上，选择或创建一个 Google Cloud 项目。

转到“项目选择器”

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

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

设置资源

本部分介绍如何为此教程设置 Cloud Storage、虚拟机和 Cloud TPU 资源。

打开一个 Cloud Shell 窗口。

打开 Cloud Shell
为项目 ID 创建一个环境变量。
```
export PROJECT_ID=project-id
```
配置 Google Cloud CLI 以使用要在其中创建 Cloud TPU 的项目。
```
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 存储桶：

注意：在以下命令中，将 bucket-name 替换为您要分配给存储桶的名称。
```
gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 gs://bucket-name
```
此 Cloud Storage 存储桶存储您用于训练模型的数据和训练结果。本教程中使用的 gcloud 命令会为您设置 TPU，还会为在上一步中设置的 Cloud TPU 服务账号设置默认权限。如果您需要更精细的权限，请查看访问级层权限。

存储桶位置必须与虚拟机 (VM) 和 TPU 节点位于同一区域。虚拟机和 TPU 节点位于特定地区，即区域内的细分。
使用 gcloud 命令启动 Compute Engine 虚拟机和 Cloud TPU。使用的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需详细了解这两种虚拟机架构，请参阅系统架构。
TPU 虚拟机
```
$ gcloud compute tpus tpu-vm create ncf-tutorial \
--zone=europe-west4-a \
--accelerator-type=v3-8 \
--version=tpu-vm-tf-2.16.1-pjrt
```
命令标志说明

zone

您计划在其中创建 Cloud TPU 的区域。

accelerator-type

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

version

Cloud TPU 软件版本。
TPU 节点
```
$ gcloud compute tpus execution-groups create  \
 --zone=europe-west4-a \
 --name=ncf-tutorial \
 --accelerator-type=v3-8 \
 --machine-type=n1-standard-8 \
 --disk-size=300 \
 --tf-version=2.12.0
```
命令标志说明

zone

您计划在其中创建 Cloud TPU 的区域。

name

TPU 名称。如果未指定，则默认为您的用户名。

accelerator-type

要创建的 Cloud TPU 的类型。

machine-type

要创建的 Compute Engine 虚拟机的机器类型。

disk-size

Compute Engine 虚拟机的根卷大小（以 GB 为单位）。

tf-version

TensorFlow gcloud 的版本会安装在虚拟机上。
注意：如果您有多个项目，则必须使用 --project 标志指定项目 ID。

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

注意：首次在项目上运行 gcloud compute tpus 时，大约需要 5 分钟来执行一些启动任务，例如 SSH 密钥传播和 API 启动。
如果您未自动登录 Compute Engine 实例，请通过运行以下 ssh 命令进行登录。登录虚拟机后，shell 提示符会从 username@projectname 更改为 username@vm-name：
TPU 虚拟机
```
gcloud compute tpus tpu-vm ssh ncf-tutorial --zone=europe-west4-a
```
TPU 节点
```
gcloud compute ssh ncf-tutorial --zone=europe-west4-a
```
要点：从现在起，前缀 (vm) $ 表示您应在 Compute Engine 虚拟机实例上运行该命令。

准备数据

为存储桶添加环境变量。将 bucket-name 替换为您的存储桶名称。
```
(vm)$ export STORAGE_BUCKET=gs://bucket-name
```

为数据目录添加环境变量。

(vm)$ export DATA_DIR=${STORAGE_BUCKET}/ncf_data

设置模型位置并设置 PYTHONPATH 环境变量。

TPU 虚拟机

(vm)$ git clone https://github.com/tensorflow/models.git
(vm)$ pip3 install -r models/official/requirements.txt

(vm)$ export PYTHONPATH="${PWD}/models:${PYTHONPATH}"

TPU 节点

(vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"
(vm)$ pip3 install -r /usr/share/models/official/requirements.txt

切换到存储模型处理文件的目录：

TPU 虚拟机

(vm)$ cd ~/models/official/recommendation

TPU 节点

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

在 DATA_DIR 中为 ml-20m 数据集生成训练和评估数据：

(vm)$ python3 create_ncf_data.py \
    --dataset ml-20m \
    --num_train_epochs 4 \
    --meta_data_file_path ${DATA_DIR}/metadata \
    --eval_prebatch_size 160000 \
    --data_dir ${DATA_DIR}

此脚本会在虚拟机上生成并预处理数据集。预处理将数据转换为模型所需的 TFRecord 格式。下载和预处理大约需要 25 分钟，并生成类似于以下内容的输出：

I0804 23:03:02.370002 139664166737728 movielens.py:124] Successfully downloaded /tmp/tmpicajrlfc/ml-20m.zip 198702078 bytes
I0804 23:04:42.665195 139664166737728 data_preprocessing.py:223] Beginning data preprocessing.
I0804 23:04:59.084554 139664166737728 data_preprocessing.py:84] Generating user_map and item_map...
I0804 23:05:20.934210 139664166737728 data_preprocessing.py:103] Sorting by user, timestamp...
I0804 23:06:39.859857 139664166737728 data_preprocessing.py:194] Writing raw data cache.
I0804 23:06:42.375952 139664166737728 data_preprocessing.py:262] Data preprocessing complete. Time: 119.7 sec.
%lt;BisectionDataConstructor(Thread-1, initial daemon)>
General:
  Num users: 138493
  Num items: 26744

Training:
  Positive count:          19861770
  Batch size:              99000
  Batch count per epoch:   1004

Eval:
  Positive count:          138493
  Batch size:              160000
  Batch count per epoch:   866

I0804 23:07:14.137242 139664166737728 data_pipeline.py:887] Negative total vector built. Time: 31.8 seconds
I0804 23:11:25.013135 139664166737728 data_pipeline.py:588] Epoch construction complete. Time: 250.9 seconds
I0804 23:15:46.391308 139664166737728 data_pipeline.py:674] Eval construction complete. Time: 261.4 seconds
I0804 23:19:54.345858 139664166737728 data_pipeline.py:588] Epoch construction complete. Time: 248.0 seconds
I0804 23:24:09.182484 139664166737728 data_pipeline.py:588] Epoch construction complete. Time: 254.8 seconds
I0804 23:28:26.224653 139664166737728 data_pipeline.py:588] Epoch construction complete. Time: 257.0 seconds

设置并开始训练 Cloud TPU

设置 Cloud TPU 名称变量。

TPU 虚拟机

(vm)$ export TPU_NAME=local

TPU 节点

(vm)$ export TPU_NAME=ncf-tutorial

运行训练和评估

以下脚本运行一个示例训练，运行三个周期：

为 Model 目录添加环境变量以保存检查点和 TensorBoard 总结：
```
(vm)$ export MODEL_DIR=${STORAGE_BUCKET}/ncf
```
创建 TPU 时，如果将 --version 参数设置为以 -pjrt 结尾的版本，请设置以下环境变量以启用 PJRT 运行时：
```
  (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
  (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
```

运行以下命令训练 NCF 模型：

(vm)$ python3 ncf_keras_main.py \
     --model_dir=${MODEL_DIR} \
     --data_dir=${DATA_DIR} \
     --train_dataset_path=${DATA_DIR}/training_cycle_*/* \
     --eval_dataset_path=${DATA_DIR}/eval_data/* \
     --input_meta_data_path=${DATA_DIR}/metadata \
     --learning_rate=3e-5 \
     --train_epochs=3 \
     --dataset=ml-20m \
     --eval_batch_size=160000 \
     --learning_rate=0.00382059 \
     --beta1=0.783529 \
     --beta2=0.909003 \
     --epsilon=1.45439e-07 \
     --dataset=ml-20m \
     --num_factors=64 \
     --hr_threshold=0.635 \
     --keras_use_ctl=true \
     --layers=256,256,128,64 \
     --use_synthetic_data=false \
     --distribution_strategy=tpu \
     --download_if_missing=false

训练和评估大约需要 2 分钟时间，并生成如下所示的最终输出：

Result is {'loss': <tf.Tensor: shape=(), dtype=float32, numpy=0.10950611>,
'train_finish_time': 1618016422.1377568, 'avg_exp_per_second': 3062557.5070816963}

清理

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

断开与 Compute Engine 实例的连接（如果您尚未这样做）：
```
(vm)$ exit
```
您的提示符现在应为 username@projectname，表明您位于 Cloud Shell 中。
删除您的 Cloud TPU 和 Compute Engine 资源。用于删除资源的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需了解详情，请参阅系统架构。
TPU 虚拟机
```
$ gcloud compute tpus tpu-vm delete ncf-tutorial \
--zone=europe-west4-a
```
TPU 节点
```
$ gcloud compute tpus execution-groups delete ncf-tutorial \
--zone=europe-west4-a
```
通过运行 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
```
```
Listed 0 items.
```
如下所示运行 gsutil，将 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 推断简介。

详细了解 ctpu，包括如何在本地机器上安装。
探索 TensorBoard 中的 TPU 工具。

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

概览

数据集

ml-1m

ml-20m

目标

费用

准备工作

设置资源

TPU 虚拟机

命令标志说明

TPU 节点

命令标志说明

TPU 虚拟机

TPU 节点

准备数据

TPU 虚拟机

TPU 节点

TPU 虚拟机

TPU 节点

设置并开始训练 Cloud TPU

TPU 虚拟机

TPU 节点

运行训练和评估

清理

TPU 虚拟机

TPU 节点

TPU 虚拟机

TPU 节点

后续步骤

超参数调优

推理