Vertex AI：利用自动打包功能，在 Vertex AI Training 上使用 Hugging Face 对 Bert 进行微调

1. 概览

在本实验中，您将学习如何使用自动打包功能在 Vertex AI Training 上运行自定义训练作业。Vertex AI 上的自定义训练作业使用容器。如果您不想构建自己的映像，可以使用自动打包功能，该功能会根据您的代码构建自定义 Docker 映像，将该映像推送到 Container Registry，并根据该映像启动 CustomJob。

学习内容

您将学习如何：

使用本地模式测试代码。
使用自动打包功能配置和启动自定义训练作业。

在 Google Cloud 上运行此实验的总费用约为 2 美元 。

2. 用例概览

您将使用 Hugging Face 中的库在 IMDB 数据集上微调 Bert 模型。该模型将预测电影评论是正面还是负面。数据集将从 Hugging Face 数据集库下载，Bert 模型将从 Hugging Face Transformers 库下载。

3. Vertex AI 简介

本实验使用的是 Google Cloud 上提供的最新 AI 产品。Vertex AI 将整个 Google Cloud 的机器学习产品集成到无缝的开发体验中。以前，使用 AutoML 训练的模型和自定义模型是通过不同的服务访问的。现在，该新产品与其他新产品一起将这两种模型合并到一个 API 中。您还可以将现有项目迁移到 Vertex AI。如果您有任何反馈，请参阅支持页面。

Vertex AI 包含许多不同的产品，可支持端到端机器学习工作流。本实验将重点介绍 Training 和 Workbench。

Vertex 产品概览

4. 设置您的环境

您需要一个启用了结算功能的 Google Cloud Platform 项目才能运行此 Codelab。如需创建项目，请按照此处的说明操作。

第 1 步：启用 Compute Engine API

前往 Compute Engine，然后选择启用（如果尚未启用）。

第 2 步：启用 Vertex AI API

前往 Cloud Console 的 Vertex AI 部分，然后点击 启用 Vertex AI API。

Vertex AI 信息中心

第 3 步：启用 Container Registry API

前往 Container Registry，然后选择启用（如果尚未启用）。您将使用此产品为您的自定义训练作业创建容器。

第 4 步：创建 Vertex AI Workbench 实例

在 Cloud Console 的 Vertex AI 部分中，点击“Workbench”：

然后点击代管式笔记本：

Notebooks_UI

然后选择新建笔记本。

new_notebook

为您的笔记本命名，然后点击高级设置。

create_notebook

在“高级设置”下，启用空闲关闭，并将分钟数设置为 60。这意味着，您的笔记本处于未使用状态时会自动关闭，以免产生不必要的费用。

idle_timeout

您可以保留所有其他高级设置。

接下来，点击创建。

创建实例后，选择打开 JupyterLab 。

open_jupyterlab

首次使用新实例时，系统会要求您进行身份验证。

authenticate

5. 编写训练代码

首先，通过“启动器”菜单在笔记本实例中打开终端窗口：

launcher_terminal

创建一个名为 autopkg-codelab 的新目录并通过 cd 命令进入该目录。

mkdir autopkg-codelab
cd autopkg-codelab

在终端中，运行以下命令来为训练代码创建一个目录，并创建一个 Python 文件，您将在其中添加代码：

mkdir trainer
touch trainer/task.py

您的 autopkg-codelab/ 目录中现在应包含以下内容：

+ trainer/
    + task.py

接下来，打开刚刚创建的 task.py 文件，然后复制以下代码。

import argparse

import tensorflow as tf
from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import TFAutoModelForSequenceClassification

CHECKPOINT = "bert-base-cased"

def get_args():
  '''Parses args.'''

  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--epochs',
      required=False,
      default=3,
      type=int,
      help='number of epochs')
  parser.add_argument(
      '--job_dir',
      required=True,
      type=str,
      help='bucket to store saved model, include gs://')
  args = parser.parse_args()
  return args


def create_datasets():
    '''Creates a tf.data.Dataset for train and evaluation.'''

    raw_datasets = load_dataset('imdb')
    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
    tokenized_datasets = raw_datasets.map((lambda examples: tokenize_function(examples, tokenizer)), batched=True)

    # To speed up training, we use only a portion of the data.
    # Use full_train_dataset and full_eval_dataset if you want to train on all the data.
    small_train_dataset = tokenized_datasets['train'].shuffle(seed=42).select(range(1000))
    small_eval_dataset = tokenized_datasets['test'].shuffle(seed=42).select(range(1000))
    full_train_dataset = tokenized_datasets['train']
    full_eval_dataset = tokenized_datasets['test']

    tf_train_dataset = small_train_dataset.remove_columns(['text']).with_format("tensorflow")
    tf_eval_dataset = small_eval_dataset.remove_columns(['text']).with_format("tensorflow")

    train_features = {x: tf_train_dataset[x] for x in tokenizer.model_input_names}
    train_tf_dataset = tf.data.Dataset.from_tensor_slices((train_features, tf_train_dataset["label"]))
    train_tf_dataset = train_tf_dataset.shuffle(len(tf_train_dataset)).batch(8)

    eval_features = {x: tf_eval_dataset[x] for x in tokenizer.model_input_names}
    eval_tf_dataset = tf.data.Dataset.from_tensor_slices((eval_features, tf_eval_dataset["label"]))
    eval_tf_dataset = eval_tf_dataset.batch(8)

    return train_tf_dataset, eval_tf_dataset


def tokenize_function(examples, tokenizer):
    '''Tokenizes text examples.'''

    return tokenizer(examples['text'], padding='max_length', truncation=True)


def main():
    args = get_args()
    train_tf_dataset, eval_tf_dataset = create_datasets()
    model = TFAutoModelForSequenceClassification.from_pretrained(CHECKPOINT, num_labels=2)

    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=tf.metrics.SparseCategoricalAccuracy(),
    )

    model.fit(train_tf_dataset, validation_data=eval_tf_dataset, epochs=args.epochs)
    model.save(f'{args.job_dir}/model_output')


if __name__ == "__main__":
    main()

关于此代码，有几点需要注意：

CHECKPOINT 是我们要微调的模型。在本例中，我们使用 Bert。
TFAutoModelForSequenceClassification 方法将在 TensorFlow 中加载指定的语言模型架构 + 权重，并在顶部添加一个具有随机初始化的权重的分类头。在本例中，我们有一个二元分类问题（正面或负面），因此我们为此分类器指定了 num_labels=2。

6. 在本地容器化和运行训练代码

您可以使用 gcloud ai custom-jobs local-run 命令根据训练代码构建 Docker 容器映像，并在本地机器上以容器的形式运行该映像。在本地运行容器采用与在 Vertex AI Training 上运行容器类似的方式执行训练代码，有助于在对 Vertex AI 执行自定义训练之前调试代码存在的问题。

在训练作业中，我们将经过训练的模型导出到 Cloud Storage 存储分区。在终端中，运行以下命令为项目定义一个环境变量，务必注意将 your-cloud-project 替换为您的项目 ID：

PROJECT_ID='your-cloud-project'

然后，创建一个存储分区。如果您有现有的存储分区，可以随意使用该存储分区。

BUCKET_NAME="gs://${PROJECT_ID}-bucket"
gsutil mb -l us-central1 $BUCKET_NAME

在 Vertex AI Training 上运行自定义训练作业时，我们将使用 GPU。但由于我们没有为 Workbench 实例指定 GPU，因此我们将使用基于 CPU 的映像进行本地测试。在本示例中，我们使用 Vertex AI Training 预构建容器。

运行以下命令来设置要用作容器基础的 Docker 映像的 URI。

BASE_CPU_IMAGE=us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-7:latest

然后，为本地运行命令构建的生成的 Docker 映像设置名称。

OUTPUT_IMAGE=$PROJECT_ID-local-package-cpu:latest

我们的训练代码使用 Hugging Face 数据集和 Transformers 库。这些库不包含在我们选择作为基础映像的映像中，因此我们需要将它们作为要求提供。为此，我们将在 autopkg-codelab 目录中创建一个 requirements.txt 文件。

确保您位于 autopkg-codelab 目录中，然后在终端中输入以下内容。

touch requirements.txt

您的 autopkg-codelab 目录中现在应包含以下内容：

+ requirements.txt
+ trainer/
    + task.py

打开要求文件，然后粘贴以下内容

datasets==1.18.2
transformers==4.16.2

最后，执行 gcloud ai custom-jobs local-run 命令，以在 Workbench 代管式实例上启动训练。

gcloud ai custom-jobs local-run \
--executor-image-uri=$BASE_CPU_IMAGE \
--python-module=trainer.task \
--output-image-uri=$OUTPUT_IMAGE \
-- \
--job_dir=$BUCKET_NAME

您应该会看到 Docker 映像正在构建。我们添加到 requirements.txt 文件中的依赖项将通过 pip 安装。首次执行此命令时，这可能需要几分钟才能完成。构建映像后，task.py 文件将开始运行，您将看到模型训练。您应会看到类似下图的界面：

local_training

由于我们未在本地使用 GPU，因此模型训练将需要很长时间。您可以按 Ctrl+c 并取消本地训练，而不是等待作业完成。

请注意，如果您想进行进一步测试，还可以直接运行上面构建的映像，而无需重新打包。

gcloud beta ai custom-jobs local-run \
--executor-image-uri=$OUTPUT_IMAGE \
-- \
--job_dir=$BUCKET_NAME \
--epochs=1

7. 创建自定义作业

现在我们已经测试了本地模式，接下来我们将使用自动打包功能在 Vertex AI Training 上启动自定义训练作业。只需一个命令，此功能即可：

根据您的代码构建自定义 Docker 映像。
将映像推送到 Container Registry。
根据映像启动 CustomJob。

返回到终端，然后 cd 到 autopkg-codelab 目录上一级。

+ autopkg-codelab
  + requirements.txt
  + trainer/
      + task.py

将 Vertex AI Training 预构建 TensorFlow GPU 映像指定为自定义训练作业的基础映像。

BASE_GPU_IMAGE=us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-7:latest

接下来，执行 gcloud ai custom-jobs create 命令。首先，此命令将根据训练代码构建自定义 Docker 映像。基础映像是我们设置为 BASE_GPU_IMAGE 的 Vertex AI Training 预构建容器。然后，自动打包功能将通过 pip 安装数据集和 Transformers 库，如我们的 requirements.txt 文件中所指定。

gcloud ai custom-jobs create \
--region=us-central1 \
--display-name=fine_tune_bert \
--args=--job_dir=$BUCKET_NAME \
--worker-pool-spec=machine-type=n1-standard-4,replica-count=1,accelerator-type=NVIDIA_TESLA_V100,executor-image-uri=$BASE_GPU_IMAGE,local-package-path=autopkg-codelab,python-module=trainer.task