XLNet:为语言理解预训练通用化自动回归 (TF 2.x)

本教程介绍如何在 Cloud TPU 上训练基于 Transformer 的双向编码器表示法 (BERT) 模型。

XLNet 是一种通用化自动回归 BERT 式预训练语言模型,可通过最大限度地提高针对因式分解顺序的所有排列的预期可能性,实现学习双向上下文。XLNet 可以通过使用 Transformer-XL 中引入的片段级重复机制和相对位置编码方案,在不中断时间连贯性的情况下获知固定时间段之外的依赖项。对于 20 个 NLP 基准任务,XLNet 的表现优于 NLP,并可获得针对 18 个任务的最佳结果,包括问答、自然语言推理、情感分析和文档排名。

如需了解详情,请参阅详细介绍 XLNet 并提供多个任务的完整结果的学术论文

目标

  • 创建 Cloud Storage 存储分区以保存数据集和模型输出。
  • 处理数据集。
  • 使用 XLNet 进行微调。
  • 验证输出结果。

费用

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

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

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

准备工作

本部分介绍如何设置 Cloud Storage 存储分区和 Compute Engine 虚拟机。

  1. 打开一个 Cloud Shell 窗口。

    打开 Cloud Shell

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

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

    gcloud config set project ${PROJECT_ID}
    

    在新的 Cloud Shell 虚拟机中首次运行此命令时,将显示“为 Cloud Shell 提供授权”页面。点击页面底部的“授权”以允许 gcloud 使用您的凭据进行 Google Cloud Platform (GCP) 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 存储分区。

    此 Cloud Storage 存储分区存储您用于训练模型的数据和训练结果。存储分区位置必须要与 Compute Engine(虚拟机)和 Cloud TPU 节点位于同一区域。

    gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 -b on gs://bucket-name
    
  6. 在使用存储分区之前,您需要向 Cloud TPU 服务帐号授予对 Cloud Storage 存储分区的访问权限。为 Cloud TPU 服务帐号设置精细 ACL

  7. 启动 Compute Engine 虚拟机和 Cloud TPU。

    $ gcloud compute tpus execution-groups create \
      --name=xlnet-tutorial \
      --machine-type=n1-standard-8 \
      --zone=europe-west4-a \
      --tf-version=2.6.0 \
      --disk-size=500GB \
      --accelerator-type=v3-8
    

    命令标志说明

    name
    要创建的 Cloud TPU 的名称。
    machine-type
    要创建的 Compute Engine 虚拟机的机器类型
    zone
    拟在其中创建 Cloud TPU 的区域
    tf-version
    要在虚拟机上使用的 TensorFlow 的版本。
    accelerator-type
    要创建的 Cloud TPU 的类型
  8. gcloud create 命令执行完毕后,验证 shell 提示符已从 username@projectname 更改为 username@vm-name。此变化表明您现已登录 Compute Engine 虚拟机。

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

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

  9. 导出 Python 目录和 Cloud Storage 存储分区名称。

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

处理数据集

数据集处理需要使用句子片段模型。可在可公开访问的 Cloud Storage 存储分区上找到模型:gs://cloud-tpu-checkpoints/xlnet/cased_spiece.model

  1. 在您的 Compute Engine 虚拟机中设置以下变量。

    (vm)$ export SPIECE_DIR=~/cased_spiece
    (vm)$ export SPIECE_MODEL=${SPIECE_DIR}/cased_spiece.model
    (vm)$ export DATASETS_DIR=${STORAGE_BUCKET}/datasets
    
    (vm)$ mkdir -p ${SPIECE_DIR}
    
    (vm)$ gsutil cp gs://cloud-tpu-checkpoints/xlnet/cased_spiece.model ${SPIECE_DIR}
    

准备分类数据集

如需在 IMDB 数据集上准备分类数据 TFRecord,请使用以下命令下载并解压缩 IMDB 数据集:

(vm)$ export IMDB_DIR=~/imdb
(vm)$ mkdir -p ${IMDB_DIR}
(vm)$ cd ${IMDB_DIR}
(vm)$ wget http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz
(vm)$ tar zxvf aclImdb_v1.tar.gz -C ${IMDB_DIR}
(vm)$ rm aclImdb_v1.tar.gz

然后,您可以使用以下命令将数据集转换为 TFRecord:

(vm)$ export TASK_NAME=imdb
(vm)$ python3 /usr/share/models/official/nlp/xlnet/preprocess_classification_data.py \
    --max_seq_length=512 \
    --spiece_model_file=${SPIECE_MODEL} \
    --output_dir=${DATASETS_DIR}/${TASK_NAME} \
    --data_dir=${IMDB_DIR}/aclImdb \
    --task_name=${TASK_NAME}

使用 XLNet 进行微调

可以在 Cloud Storage 文件夹 gs://cloud-tpu-checkpoints/xlnet/keras_xlnet 中找到解压缩的预训练模型文件。

  1. 导出预训练模型和输出 Cloud Storage 存储分区的变量。

    (vm)$ export XLNET_DIR=gs://cloud-tpu-checkpoints/xlnet/keras_xlnet
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/xlnet-output
    

训练模型

以下示例代码会在 IMDB 数据集上优化 XLNet。对于此任务,获取前 500 个步骤的结果大约需要 11 分钟,在 v3-8 中完成大约需要 1 小时。您应该会看到准确率介于 96.15 和 96.33 之间。

(vm)$ export TPU_NAME=xlnet-tutorial
(vm)$ python3 /usr/share/models/official/nlp/xlnet/run_classifier.py \
  --strategy_type=tpu \
  --tpu=${TPU_NAME} \
  --init_checkpoint=${XLNET_DIR}/xlnet_model.ckpt \
  --model_dir=${MODEL_DIR} \
  --test_data_size=25024 \
  --train_tfrecord_path=${DATASETS_DIR}/imdb/cased_spiece.model.len-512.train.tf_record \
  --test_tfrecord_path=${DATASETS_DIR}/imdb/cased_spiece.model.len-512.dev.eval.tf_record \
  --train_batch_size=32 \
  --seq_len=512 \
  --n_layer=24 \
  --d_model=1024 \
  --d_embed=1024 \
  --n_head=16 \
  --d_head=64 \
  --d_inner=4096 \
  --untie_r=true \
  --n_class=2 \
  --ff_activation=gelu \
  --learning_rate=2e-5 \
  --train_steps=4000 \
  --warmup_steps=500 \
  --iterations=500 \
  --bi_data=false \
  --summary_type=last

命令标志说明

mode
设为 train_and_eval 时,此脚本用于训练和评估模型。如果设置为“export_only”,此脚本会导出保存的模型。
input_meta_data_path
文件的路径,该文件包含要用于训练和评估的数据集的相关元数据。
train_data_path
训练输入的 Cloud Storage 路径。在此示例中,该路径设置为 fake_imagenet 数据集。
eval_data_path
评估输入的 Cloud Storage 路径。在此示例中,该路径设置为 fake_imagenet 数据集。
init_checkpoint
包含预训练 BERT 模型的初始检查点的 json 文件路径。
train_batch_size
训练批次大小。
eval_batch_size
评估批次大小。
learning_rate
学习速率。
num_train_epochs
使用整个数据集训练模型的次数。
model_dir
用于指定在模型训练期间存储检查点和摘要的目录。如果指定的文件夹不存在,此程序会自行创建。使用 Cloud TPU 时,model_dir 必须是 Cloud Storage 路径 (`gs://...`)。您可以重复使用现有的文件夹来加载当前检查点数据和存储其他检查点,只要先前的检查点是使用相同大小的 TPU 和相同 Tensorflow 版本创建的即可。
distribution_strategy
要在 TPU 上训练 ResNet 模型,您必须将 distribution_strategy 设置为 tpu
tpu
Cloud TPU 的名称。这通过指定环境变量 (TPU_NAME) 进行设置。

验证结果

此训练在 v3-8 TPU 上大约需要 1 小时。 脚本完成后,您应该会看到如下所示的结果:

Training Summary:
{'total_training_steps': 4000, 'train_loss': 0.00868087075650692,
'last_train_metrics': 0.9983749985694885, 'eval_metrics': 0.96148}

清理

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

    (vm)$ exit
    

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

  2. 在 Cloud Shell 中,运行以下 gcloud 命令以删除您的虚拟机和 Cloud TPU。

    $ gcloud compute tpus execution-groups delete xlnet-tutorial \
      --zone=europe-west4-a
    
  3. 运行 gcloud compute tpus list 以确保未分配任何 Cloud TPU 实例,避免产生不必要的 TPU 使用费。删除操作可能需要几分钟时间才能完成。

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

    该命令后面的输出列表中不应显示您使用的 Cloud TPU:

    NAME            ZONE            ACCELERATOR_TYPE  NETWORK_ENDPOINTS   NETWORK  RANGE            STATUS  API_VERSION
    xyz-tutorial    europe-west4-a  v3-8              10.119.91.74:8470   default  10.119.91.72/29  READY   V1_ALPHA1
    
  4. 如下所示运行 gsutil,将 bucket-name 替换为您为本教程创建的 Cloud Storage 存储分区的名称:

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

后续步骤

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

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

超参数调节

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

推理

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