此页面由 Cloud Translation API 翻译。

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 数据集库下载，并从 Hugging Face 转换器库中下载 Bert 模型。

3. Vertex AI 简介

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

Vertex AI 包含许多不同的产品，可支持端到端机器学习工作流。本实验将重点介绍训练和 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

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

身份验证

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。但是，由于我们没有指定带有 GPU 的 Workbench 实例，因此我们将使用基于 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 数据集和转换器库。这些库未包含在我们选择作为基础映像的映像中，因此我们需要将其作为要求提供。为此，我们将在 autopkg-codelab 目录中创建一个 requirements.txt 文件。

确保您位于 autopkg-codelab 目录中，并在终端中输入以下命令。

touch requirements.txt

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

+ requirements.txt
+ trainer/
    + task.py

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

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 映像。系统将通过 pip 安装我们添加到 requirements.txt 文件的依赖项。首次执行此命令时，可能需要几分钟才能完成。构建映像后，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 预构建容器。然后，自动打包功能将按照 requirements.txt 文件中指定的方式使用 pip 安装数据集和转换器库。

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

我们来看一下 worker-pool-spec 实参。这定义了自定义作业使用的工作器池配置。您可以指定多个工作器池规范，以创建具有多个工作器池的自定义作业，以进行分布式训练。在此示例中，我们仅指定一个工作器池，因为我们的训练代码未针对分布式训练进行配置。

以下是本规范的一些关键字段：