Vertex AI：分布式超参数调节

1. 概览

在本实验中，您将学习如何使用 Vertex AI 进行超参数调节和分布式训练。虽然本实验使用 TensorFlow 构建模型代码，但相关概念也适用于其他机器学习框架。

学习内容

您将了解如何：

  • 在自定义容器中使用分布式训练来训练模型
  • 启动训练代码的多次试验以自动进行超参数调节

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

2. Vertex AI 简介

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

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

Vertex 产品概览

3. 用例概览

在本实验中，您将使用超参数调节来发现通过 TensorFlow 数据集中的马或人数据集训练的图片分类模型的最佳参数。

超参数调节

使用 Vertex AI Training 进行超参数调节的工作方式是，使用您所选超参数的值（设置在您指定的限制范围内）运行训练应用的多次试验。Vertex AI 会跟踪每次试验的结果，并针对后续试验进行调整。

如需将超参数调节与 Vertex AI Training 结合使用，您需要对训练代码进行以下两项更改：

  1. 在主训练模块中，为您要调节的每个超参数定义一个命令行参数。
  2. 使用这些参数中传递的值，在应用的代码中设置相应的超参数。

分布式训练

如果您有一个 GPU，则 TensorFlow 会使用此加速器加快模型训练速度，您无需执行额外的操作。但是，如果您希望使用多个 GPU 进一步提升速度，则需要使用 tf.distribute，它是 TensorFlow 的模块，用于在多个设备上运行计算。

本实验使用 tf.distribute.MirroredStrategy，您只需对代码进行一些更改，即可将其添加到训练应用中。此策略会在机器的每个 GPU 上创建模型的副本。后续的梯度更新会以同步的方式进行。这意味着，每个 GPU 都会针对不同的输入数据切片计算通过模型的前向和后向传递。然后，来自各个切片的计算梯度会在所有 GPU 上聚合，并在一个称为 all-reduce 的进程中取平均值。模型参数会使用这些平均梯度进行更新。

您无需了解相关详细信息即可完成本实验，但如果您想详细了解分布式训练在 TensorFlow 中的工作原理，请观看以下视频：

4. 设置您的环境

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

第 1 步：启用 Compute Engine API

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

第 2 步：启用 Container Registry API

前往 Container 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

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

idle_timeout

在“安全性”下，选择“启用终端”（如果尚未启用）。

enable-terminal

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

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

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

open_jupyterlab

首次使用新实例时，系统会要求您进行身份验证。为此，请按照界面中的步骤操作。

身份验证

5. 编写训练代码

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

launcher_terminal

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

mkdir vertex-codelab
cd vertex-codelab

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

mkdir trainer
touch trainer/task.py

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

+ trainer/
    + task.py

接下来，打开您刚刚创建的 task.py 文件，并在其中粘贴下面的所有代码。

import tensorflow as tf
import tensorflow_datasets as tfds
import argparse
import hypertune
import os

NUM_EPOCHS = 10
BATCH_SIZE = 64

def get_args():
  '''Parses args. Must include all hyperparameters you want to tune.'''

  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--learning_rate',
      required=True,
      type=float,
      help='learning rate')
  parser.add_argument(
      '--momentum',
      required=True,
      type=float,
      help='SGD momentum value')
  parser.add_argument(
      '--num_units',
      required=True,
      type=int,
      help='number of units in last hidden layer')
  args = parser.parse_args()
  return args

def preprocess_data(image, label):
  '''Resizes and scales images.'''

  image = tf.image.resize(image, (150,150))
  return tf.cast(image, tf.float32) / 255., label

def create_dataset(batch_size):
  '''Loads Horses Or Humans dataset and preprocesses data.'''

  data, info = tfds.load(name='horses_or_humans', as_supervised=True, with_info=True)

  # Create train dataset
  train_data = data['train'].map(preprocess_data)
  train_data  = train_data.shuffle(1000)
  train_data  = train_data.batch(batch_size)

  # Create validation dataset
  validation_data = data['test'].map(preprocess_data)
  validation_data  = validation_data.batch(batch_size)

  return train_data, validation_data

def create_model(num_units, learning_rate, momentum):
  '''Defines and compiles model.'''

  inputs = tf.keras.Input(shape=(150, 150, 3))
  x = tf.keras.layers.Conv2D(16, (3, 3), activation='relu')(inputs)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Conv2D(32, (3, 3), activation='relu')(x)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Conv2D(64, (3, 3), activation='relu')(x)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Flatten()(x)
  x = tf.keras.layers.Dense(num_units, activation='relu')(x)
  outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)
  model = tf.keras.Model(inputs, outputs)
  model.compile(
      loss='binary_crossentropy',
      optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=momentum),
      metrics=['accuracy'])
  return model

def main():
  args = get_args()

  # Create distribution strategy
  strategy = tf.distribute.MirroredStrategy()

  # Get data
  GLOBAL_BATCH_SIZE = BATCH_SIZE * strategy.num_replicas_in_sync
  train_data, validation_data = create_dataset(GLOBAL_BATCH_SIZE)

  # Wrap variable creation within strategy scope
  with strategy.scope():
    model = create_model(args.num_units, args.learning_rate, args.momentum)

  # Train model
  history = model.fit(train_data, epochs=NUM_EPOCHS, validation_data=validation_data)

  # Define metric
  hp_metric = history.history['val_accuracy'][-1]

  hpt = hypertune.HyperTune()
  hpt.report_hyperparameter_tuning_metric(
      hyperparameter_metric_tag='accuracy',
      metric_value=hp_metric,
      global_step=NUM_EPOCHS)

if __name__ == "__main__":
    main()

我们来深入了解一下代码，并检查特定于分布式训练和超参数调节的组成部分。

分布式训练

  1. 在 main() 函数中，相关代码创建了 MirroredStrategy 对象。接下来，您将模型变量的创建过程封装在策略的范围内。此步骤指示 TensorFlow 应在 GPU 中镜像哪些变量。
  2. 批次大小按 num_replicas_in_sync 扩容。在 TensorFlow 中使用同步数据并行策略时，扩缩批次大小是最佳实践。您可以点击此处了解详情。

超参数调节

  1. 该脚本会导入 hypertune 库。之后，在构建容器映像时，我们需要确保安装此库。
  2. get_args() 函数为您要调节的每个超参数定义命令行参数。在此示例中，将要调节的超参数是学习速率、优化器中的动量值以及模型最后一个隐藏层中的单元数，但您可以随意尝试其他超参数。然后，在这些参数中传递的值将用于在代码中设置相应的超参数（例如，设置 learning_rate = args.learning_rate）。
  3. 在 main() 函数的末尾，hypertune 库用于定义您要优化的指标。在 TensorFlow 中，Keras model.fit 方法会返回 History 对象。History.history 特性是连续周期的训练损失值和指标值的记录。如果您将验证数据传递给 model.fit，则 History.history 特性也将包含验证损失和指标值。例如，如果您使用验证数据将模型训练三个周期并提供 accuracy 作为指标，则 History.history 特性将类似于以下字典。
{
 "accuracy": [
   0.7795261740684509,
   0.9471358060836792,
   0.9870933294296265
 ],
 "loss": [
   0.6340447664260864,
   0.16712145507335663,
   0.04546636343002319
 ],
 "val_accuracy": [
   0.3795261740684509,
   0.4471358060836792,
   0.4870933294296265
 ],
 "val_loss": [
   2.044623374938965,
   4.100203514099121,
   3.0728273391723633
 ]

如果您希望超参数调节服务发现使模型的验证准确率最大化的值，则可以将指标定义为 val_accuracy 列表的最后一个条目（或 NUM_EPOCS - 1）。然后，将此指标传递给 HyperTune 的实例。您可以为 hyperparameter_metric_tag 选择任何字符串，但稍后在启动超参数调节作业时，您需要再次使用该字符串。

6. 容器化代码

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

第 1 步：编写 Dockerfile

在终端中，确保当前目录为 vertex-codelab，然后创建一个空的 Dockerfile：

touch Dockerfile

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

+ Dockerfile
+ trainer/
    + task.py

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

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

WORKDIR /

# Installs hypertune library
RUN pip install cloudml-hypertune

# 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"]

此 Dockerfile 使用 Deep Learning Container TensorFlow Enterprise 2.7 GPU Docker 映像。Google Cloud 上的 Deep Learning Containers 预安装了许多常见的机器学习和数据科学框架。下载该映像后，此 Dockerfile 会为训练代码设置入口点。

第 2 步：构建容器

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

PROJECT_ID='your-cloud-project'

使用 Google Container Registry 中容器映像的 URI 定义变量：

IMAGE_URI="gcr.io/$PROJECT_ID/horse-human-codelab:latest"

配置 Docker

gcloud auth configure-docker

然后，从 vertex-codelab 目录的根目录运行以下命令来构建容器：

docker build ./ -t $IMAGE_URI

最后，将其推送到 Google Container Registry：

docker push $IMAGE_URI

第 3 步：创建 Cloud Storage 存储桶

在训练作业中，我们将传入暂存存储桶的路径。

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

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

7. 启动超参数调节作业

第 1 步：使用超参数调节创建自定义训练作业

在启动器中，打开一个新的 TensorFlow 2 笔记本。

new_notebook

导入 Vertex AI Python SDK。

from google.cloud import aiplatform
from google.cloud.aiplatform import hyperparameter_tuning as hpt

如需启动超参数调节作业，您需要先定义用于指定机器类型和 Docker 映像的 worker_pool_specs。以下规范定义了一台具有两个 NVIDIA Tesla V100 GPU 的机器。

您需要将 image_uri 中的 {PROJECT_ID} 替换为您的项目。

# The spec of the worker pools including machine type and Docker image
# Be sure to replace PROJECT_ID in the "image_uri" with your project.

worker_pool_specs = [{
    "machine_spec": {
        "machine_type": "n1-standard-4",
        "accelerator_type": "NVIDIA_TESLA_V100",
        "accelerator_count": 2
    },
    "replica_count": 1,
    "container_spec": {
        "image_uri": "gcr.io/{PROJECT_ID}/horse-human-codelab:latest"
    }
}]

接下来，定义 parameter_spec，这是一个字典，用于指定您要优化的参数。字典键是您分配给每个超参数的命令行参数的字符串，字典值是参数规范。

对于每个超参数，您需要定义类型以及调节服务将尝试的值的边界。超参数可以是双精度、整数、分类或离散类型。如果您选择双精度或整数类型，则需要提供最小值和最大值。如果您选择分类或离散类型，则需要提供值。对于双精度和整数类型，您还需要提供调节值。您可以通过此视频详细了解如何选择最佳调节值。

# Dictionary representing parameters to optimize.
# The dictionary key is the parameter_id, which is passed into your training
# job as a command line argument,
# And the dictionary value is the parameter specification of the metric.
parameter_spec = {
    "learning_rate": hpt.DoubleParameterSpec(min=0.001, max=1, scale="log"),
    "momentum": hpt.DoubleParameterSpec(min=0, max=1, scale="linear"),
    "num_units": hpt.DiscreteParameterSpec(values=[64, 128, 512], scale=None)
}

要定义的最终规范是 metric_spec，它是表示要优化的指标的字典。字典键是您在训练应用代码中设置的 hyperparameter_metric_tag，值是优化目标。

# Dicionary representing metrics to optimize.
# The dictionary key is the metric_id, which is reported by your training job,
# And the dictionary value is the optimization goal of the metric.
metric_spec={'accuracy':'maximize'}

定义这些规范后，您将创建一个 CustomJob，这是将要用于在每次超参数调节试验中运行作业的常见规范。

您需要将 {YOUR_BUCKET} 替换为之前创建的存储桶。

# Replace YOUR_BUCKET
my_custom_job = aiplatform.CustomJob(display_name='horses-humans',
                              worker_pool_specs=worker_pool_specs,
                              staging_bucket='gs://{YOUR_BUCKET}')

然后，创建并运行 HyperparameterTuningJob。

hp_job = aiplatform.HyperparameterTuningJob(
    display_name='horses-humans',
    custom_job=my_custom_job,
    metric_spec=metric_spec,
    parameter_spec=parameter_spec,
    max_trial_count=6,
    parallel_trial_count=2,
    search_algorithm=None)

hp_job.run()

请注意以下几个参数：

  • max_trial_count：您需要对服务将运行的试验次数设置上限。更多的试验通常会带来更好的结果，但会有一个收益递减点，在该点之后，额外的试验对您尝试优化的指标几乎没有什么效果。最佳实践是从少量试验开始，并在扩容之前了解您所选超参数的影响力。
  • parallel_trial_count：如果您使用并行试验，则服务会预配多个训练处理集群。增加并行试验的次数会减少超参数调节作业运行所需的时间，但可能降低作业的整体效率。这是因为默认调节策略使用先前试验的结果来告知后续试验中的赋值情况。
  • search_algorithm：您可以将搜索算法设置为网格、随机或默认（无）。默认选项会应用贝叶斯优化来搜索可能的超参数值空间，并且是推荐的算法。如需详细了解此算法，请点击此处。

作业开始后，您将能够在界面中的超参数调节作业标签页下跟踪状态。

HP_job

作业完成后，您可以查看试验结果并对其进行排序，以发现超参数值的最佳组合。

HP_results

🎉 恭喜！🎉

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

  • 使用分布式训练运行超参数调节作业

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

8. 清理

因为我们将笔记本配置为在空闲 60 分钟后超时，所以不必担心关停实例。如果您要手动关停实例，请点击控制台的 Vertex AI Workbench 部分中的“停止”按钮。如果您想完全删除该笔记本，请点击“删除”按钮。

删除

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

删除存储空间