在生产环境中进行原型设计:使用 Vertex AI 训练自定义模型

1. 概览

在本实验中,您将使用 Vertex AI 运行一个自定义训练作业。

本实验是对生产环境进行原型设计视频系列的一部分。您将使用 Flowers 数据集构建图片分类模型。您可以观看随附的视频了解详情:

学习内容

您将了解如何:

  • 创建 Vertex AI Workbench 代管式笔记本
  • 从 Vertex AI 界面配置并启动自定义训练作业
  • 使用 Vertex AI Python SDK 配置并启动自定义训练作业

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

2. Vertex AI 简介

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

Vertex AI 包含许多不同的产品,可支持端到端机器学习工作流。本实验将重点介绍下面突出显示的产品:TrainingWorkbench

Vertex 产品概览

3. 设置您的环境

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

第 1 步:启用 Compute Engine API

前往 Compute Engine,然后选择启用(如果尚未启用)。

第 2 步:启用 Artifact Registry API

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

第 3 步:启用 Vertex AI API

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

Vertex AI 信息中心

第 4 步:创建 Vertex AI Workbench 实例

在 Cloud Console 的 Vertex AI 部分中,点击“Workbench”:

Vertex AI 菜单

启用 Notebooks API(如果尚未启用)。

Notebook_api

启用后,点击代管式笔记本

Notebooks_UI

然后选择新建笔记本

new_notebook

为您的笔记本命名,然后在权限下选择服务帐号

create_notebook

选择高级设置

安全性下,选择“启用终端”(如果尚未启用)。

enable_terminal

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

接下来,点击创建。预配实例需要几分钟时间。

创建实例后,选择打开 JUPYTERLAB

open_jupyterlab

4. 容器化训练应用代码

如需将此训练作业提交到 Vertex AI,您可以将训练应用代码放在 Docker 容器中,然后将此容器推送到 Google Artifact Registry。使用此方法,您可以训练使用任何框架构建的模型。

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

在笔记本中打开终端

第 1 步:创建 Cloud Storage 存储桶

在此训练作业中,您需要将经过训练的 TensorFlow 模型导出到 Cloud Storage 存储桶。您还需要将用于训练的数据存储在 Cloud Storage 存储桶中。

在终端中,运行以下命令为项目定义一个环境变量,务必注意将 your-cloud-project 替换为您的项目 ID:

PROJECT_ID='your-cloud-project'

接下来,在终端中运行以下命令,以便在项目中创建一个新的存储桶。

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

第 2 步:将数据复制到 Cloud Storage 存储桶

我们需要将花卉数据集放入 Cloud Storage 中。出于演示目的,您需要先将数据集下载到此 Workbench 实例,然后将其复制到存储桶。

下载并解压缩数据。

wget https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz
tar xvzf flower_photos.tgz

然后,将其复制到您刚创建的存储桶。我们之所以添加 -r,是因为我们想要复制整个目录,而添加 -m 则可以执行多处理副本,从而加快处理速度。

gsutil -m cp -r flower_photos $BUCKET

第 3 步:编写训练代码

创建一个名为 flowers 的新目录并通过 cd 命令进入该目录:

mkdir flowers
cd flowers

运行以下命令,为训练代码创建一个目录和一个 Python 文件(您将在其中添加代码)。

mkdir trainer
touch trainer/task.py

您的 flowers/ 目录中现在应包含以下内容:

+ trainer/
    + task.py

如需详细了解如何设计训练应用代码的结构,请参阅相关文档

接下来,打开您刚刚创建的 task.py 文件并复制下面的代码。

您需要将 {your-gcs-bucket} 替换为您刚刚创建的 Cloud Storage 存储桶的名称。

通过 Cloud Storage FUSE 工具,Vertex AI Training 上的训练作业可以将 Cloud Storage 上的数据作为本地文件系统中的文件进行访问。当您启动自定义训练作业时,该作业会看到一个目录 /gcs,该目录包含所有 Cloud Storage 存储桶(作为子目录)。因此,训练代码中的数据路径以 /gcs 开头。

import tensorflow as tf
import numpy as np
import os

## Replace {your-gcs-bucket} !!
BUCKET_ROOT='/gcs/{your-gcs-bucket}'

# Define variables
NUM_CLASSES = 5
EPOCHS=10
BATCH_SIZE = 32

IMG_HEIGHT = 180
IMG_WIDTH = 180

DATA_DIR = f'{BUCKET_ROOT}/flower_photos'

def create_datasets(data_dir, batch_size):
  '''Creates train and validation datasets.'''

  train_dataset = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="training",
    seed=123,
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    batch_size=batch_size)

  validation_dataset = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="validation",
    seed=123,
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    batch_size=batch_size)

  train_dataset = train_dataset.cache().shuffle(1000).prefetch(buffer_size=tf.data.AUTOTUNE)
  validation_dataset = validation_dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)

  return train_dataset, validation_dataset

def create_model():
  '''Creates model.'''

  model = tf.keras.Sequential([
    tf.keras.layers.Resizing(IMG_HEIGHT, IMG_WIDTH),
    tf.keras.layers.Rescaling(1./255, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3)),
    tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')
  ])
  return model

# CREATE DATASETS
train_dataset, validation_dataset = create_datasets(DATA_DIR, BATCH_SIZE)

# CREATE/COMPILE MODEL
model = create_model()
model.compile(optimizer=tf.keras.optimizers.Adam(),
              loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])

# TRAIN MODEL
history = model.fit(
  train_dataset,
  validation_data=validation_dataset,
  epochs=EPOCHS
)

# SAVE MODEL
model.save(f'{BUCKET_ROOT}/model_output')

第 4 步:创建 Dockerfile

如需将代码容器化,您需要创建一个 Dockerfile。在 Dockerfile 中,您将添加运行映像所需的所有命令。它将安装所有必要的库并为训练代码设置入口点。

通过终端,在 flowers 目录的根目录下创建一个空的 Dockerfile:

touch Dockerfile

您的 flowers/ 目录中现在应包含以下内容:

+ Dockerfile
+ trainer/
    + task.py

打开 Dockerfile 并将以下代码复制到其中:

FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-8

WORKDIR /

# Copies the trainer code to the docker image.
COPY trainer /trainer

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

我们来看看此文件中的命令。

FROM 命令可指定基础映像,基础映像是创建映像所基于的父映像。作为基础映像,您将使用 Deep Learning Container TensorFlow Enterprise 2.8 GPU Docker 映像。Google Cloud 上的 Deep Learning Containers 预安装了许多常见的机器学习和数据科学框架。

WORKDIR 命令用于指定映像上运行后续指令的目录。

COPY 命令将训练程序代码复制到 Docker 映像。请注意,在此示例中,我们的训练程序目录中只有一个 Python 文件,但对于更实际的示例,您可能还有其他文件。也许有一个名为 data.py 的文件(该文件对数据进行预处理),以及一个名为 model.py 的文件(该文件只有模型代码),等等。如需了解更复杂的训练代码,请查看有关对 Python 项目进行打包的 Python 文档。

如果您想要添加任何其他库,可以使用 RUN 命令通过 pip 安装(例如:RUN pip install -r requirements.txt)。但对于此处的示例,我们不需要任何其他内容。

最后,ENTRYPOINT 命令会设置入口点以调用训练程序。这是启动训练作业时将运行的内容。本例中执行的是 task.py 文件。

您可以在此处详细了解如何为 Vertex AI Training 编写 Dockerfile。

第 4 步:构建容器

在您的 Workbench 笔记本的终端中,运行以下命令为项目定义一个环境变量,务必将 your-cloud-project 替换为您的项目 ID:

PROJECT_ID='your-cloud-project'

在 Artifact Registry 中创建代码库

REPO_NAME='flower-app'

gcloud artifacts repositories create $REPO_NAME --repository-format=docker \
--location=us-central1 --description="Docker repository"

使用 Google Artifact Registry 中容器映像的 URI 定义变量:

IMAGE_URI=us-central1-docker.pkg.dev/$PROJECT_ID/$REPO_NAME/flower_image:latest

配置 Docker

gcloud auth configure-docker \
    us-central1-docker.pkg.dev

然后,从 flower 目录的根目录运行以下命令来构建容器:

docker build ./ -t $IMAGE_URI

最后,将其推送到 Artifact Registry:

docker push $IMAGE_URI

在将容器推送到 Artifact Registry 后,您就可以启动训练作业了。

5. 在 Vertex AI 上运行自定义训练作业

本实验通过 Google Artifact Registry 上的自定义容器使用自定义训练,但您也可以使用预构建容器运行训练作业。

首先,请前往 Cloud 控制台的“Vertex”部分中的训练部分:

“训练”菜单

第 1 步:配置训练作业

点击创建以输入训练作业的参数。

create_training

  • 数据集下,选择无代管式数据集
  • 然后选择自定义训练(高级)作为训练方法,并点击继续
  • 选择训练新模型,然后输入 flowers-model(或者用于指代模型的任何名称)作为模型名称
  • 点击 Continue(继续)

在“容器设置”步骤中,选择自定义容器

自定义容器选项

在第一个框(容器映像)中,输入上一部分中的 IMAGE_URI 变量的值。它应该为 us-central1-docker.pkg.dev/{PROJECT_ID}/flower-app/flower_image:latest(包含您自己的项目 ID)。将其余字段留空,然后点击继续

再次点击继续,跳过超参数步骤。

第 2 步:配置计算集群

按如下方式配置工作器池 0

worker_pool_0

您将暂时跳过第 6 步,并在本系列的下一个实验中配置预测容器。

点击开始训练以启动训练作业。在控制台的“训练”部分中的训练流水线标签页下,您会看到刚启动的作业:

训练作业

🎉 恭喜!🎉

您学习了如何使用 Vertex AI 执行以下操作:

  • 启动自定义训练作业,以处理自定义容器中提供的训练代码。您在本示例中使用了 TensorFlow 模型,但可以使用自定义或内置容器训练通过任何框架构建的模型。

如需详细了解 Vertex 的不同部分,请参阅相关文档

6. [可选] 使用 Vertex AI Python SDK

上一部分介绍了如何通过界面启动训练作业。在本部分中,您会看到另一种使用 Vertex AI Python SDK 提交训练作业的方法。

返回您的笔记本实例,并通过启动器创建 TensorFlow 2 笔记本:

new_notebook

导入 Vertex AI SDK。

from google.cloud import aiplatform

然后,创建一个 CustomContainerTrainingJob。您需要将 container_uri 中的 {PROJECT_ID} 替换为您的项目的名称,并将 staging_bucket 中的 {BUCKET} 替换为您之前创建的存储桶。

my_job = aiplatform.CustomContainerTrainingJob(display_name='flower-sdk-job',
                                               container_uri='us-central1-docker.pkg.dev/{PROJECT_ID}/flower-app/flower_image:latest',
                                               staging_bucket='gs://{BUCKET}')

然后,运行该作业。

my_job.run(replica_count=1,
           machine_type='n1-standard-8',
           accelerator_type='NVIDIA_TESLA_V100',
           accelerator_count=1)

为便于演示,该作业已配置为在比上一部分中的机器更大的机器上运行。此外,我们还要使用 GPU 运行。如果您不指定 machine-typeaccelerator_typeaccelerator_count,则作业默认在 n1-standard-4 上运行。

在控制台的“训练”部分中的自定义作业标签页下,您将看到您的训练作业。

7. 清理

由于 Vertex AI Workbench 代管式笔记本具有空闲关停功能,因此我们无需操心关停实例的事。如果您要手动关停实例,请点击控制台的 Vertex AI Workbench 部分中的“停止”按钮。如果您想完全删除该笔记本,请点击“删除”按钮。

可停止实例

如需删除存储桶,请使用 Cloud Console 中的导航菜单,浏览到“存储空间”,选择您的存储桶,然后点击“删除”:

删除存储空间