在云端使用 GPU 训练模型

图形处理单元 (GPU) 可显著加快许多深度学习模型的训练过程。用于图片分类、视频分析和自然语言处理等任务的训练模型涉及计算密集型矩阵乘法以及其他可利用 GPU 大规模并行架构的操作。

对于需要对超大数据集执行密集计算任务的深度学习模型,可能需要在单个处理器上运行数日才能完成训练。但是,如果您将程序设计为将这些任务分流到一个或多个 GPU,则可以将训练时间从数日缩短至数小时。

准备工作

通过 AI Platform Training,您可以在支持 GPU 的机器上运行 TensorFlow 训练应用。请参阅有关使用 GPU 的 TensorFlow 指南本文档中有关调整训练代码以利用 GPU 的部分,以确保您的应用使用可用的 GPU。

如果您使用自定义容器进行训练,还可以将 GPU 与 TensorFlow 以外的机器学习框架搭配使用。

在 GPU 上运行并不能让某些模型受益。我们建议您使用 GPU 来执行需要完成许多数学运算的大型复杂模型。即便如此,您也应该通过训练来运行较小的数据示例,以测试 GPU 支持的优势。

请求支持 GPU 的机器

要在云端使用 GPU,请通过以下任一方式将训练作业配置为可以访问支持 GPU 的机器:

  • 使用 BASIC_GPU 规模层级。
  • 使用 Compute Engine 机器类型并挂接 GPU。
  • 使用支持 GPU 的旧版机器类型。

支持基本 GPU 的机器

如果您要学习如何使用 AI Platform Training 或对支持 GPU 的机器进行实验,则可以将规模层级设置为 BASIC_GPU,以获取配备一个 NVIDIA Tesla K80 GPU 的单一工作器实例。

挂接了 GPU 的 Compute Engine 机器类型

如果您将训练作业配置为使用 Compute Engine 机器类型,则可以挂接自定义数量的 GPU 以加快作业速度。

  • 将规模层级设置为 CUSTOM
  • 将您的主实例工作器和作业中的任何其他任务类型(工作器、参数服务器或评估程序)配置为使用有效的 Compute Engine 机器类型。
  • 添加 acceleratorConfig 字段,其中包含您希望 masterConfigworkerConfigparameterServerConfigevaluatorConfig 所需的 GPU 类型和数量,具体取决于您想要加速的虚拟机 (VM) 实例。您可以使用以下 GPU 类型:
    • NVIDIA_TESLA_A100
    • NVIDIA_TESLA_K80
    • NVIDIA_TESLA_P4
    • NVIDIA_TESLA_P100
    • NVIDIA_TESLA_T4
    • NVIDIA_TESLA_V100

要创建有效的 acceleratorConfig,您必须考虑以下几项限制:

  1. 您的配置只能使用特定数量的 GPU。例如,您可以挂接 2 个或 4 个 NVIDIA Tesla T4,但不能挂接 3 个。要了解对每种类型的 GPU 有效的限制数量,请参阅下面的兼容性表

  2. 您必须确保每项 GPU 配置都会为您挂接了 GPU 的机器类型提供足够的虚拟 CPU 和内存。例如,如果您对工作器使用 n1-standard-32,则每个工作器将配备 32 个虚拟 CPU 和 120 GB 的内存。由于每个 NVIDIA Tesla V100 最多提供 12 个虚拟 CPU 和 76 GB 的内存,因此您必须至少将 4 个 NVIDIA Tesla V100 挂接到每个 n1-standard-32 工作器才能满足其需求。(2 个 GPU 提供的资源就已足够,您不能指定 3 个 GPU。)

    查看 AI Platform Training 的机器类型列表比较用于处理计算工作负载的 GPU 以确定这些兼容性,或参考下面的兼容性表

    请注意,以下是在特定情况下针对 AI Platform Training 的 GPU 资源的其他几项限制:

    • 配备了 8 个 NVIDIA Tesla K80 GPU 的配置在所有区域和地区最多只提供 208 GB 的内存。
    • 配备了 4 个 NVIDIA Tesla P100 GPU 的配置在所有区域和地区最多只支持 64 个虚拟 CPU 和 208 GB 的内存。
  3. 您必须将训练作业提交到支持 GPU 配置的区域。参阅下面的区域支持

下表提供了可以挂接到各 Compute Engine 机器类型的各类加速器数量的快速参考信息:

每种机器类型的有效 GPU 数量
机器类型 NVIDIA A100 NVIDIA Tesla K80 NVIDIA Tesla P4 NVIDIA Tesla P100 NVIDIA Tesla T4 NVIDIA Tesla V100
n1-standard-4 1、2、4、8 1、2、4 1、2、4 1、2、4 1、2、4、8
n1-standard-8 1、2、4、8 1、2、4 1、2、4 1、2、4 1、2、4、8
n1-standard-16 2、4、8 1、2、4 1、2、4 1、2、4 2、4、8
n1-standard-32 4、8 2、4 2、4 2、4 4、8
n1-standard-64 4 4 8
n1-standard-96 4 4 8
n1-highmem-2 1、2、4、8 1、2、4 1、2、4 1、2、4 1、2、4、8
n1-highmem-4 1、2、4、8 1、2、4 1、2、4 1、2、4 1、2、4、8
n1-highmem-8 1、2、4、8 1、2、4 1、2、4 1、2、4 1、2、4、8
n1-highmem-16 2、4、8 1、2、4 1、2、4 1、2、4 2、4、8
n1-highmem-32 4、8 2、4 2、4 2、4 4、8
n1-highmem-64 4 4 8
n1-highmem-96 4 4 8
n1-highcpu-16 2、4、8 1、2、4 1、2、4 1、2、4 2、4、8
n1-highcpu-32 4、8 2、4 2、4 2、4 4、8
n1-highcpu-64 8 4 4 4 8
n1-highcpu-96 4 4 8
a2-highgpu-1g 1
a2-highgpu-2g 2
a2-highgpu-4g 4
a2-highgpu-8g 8
a2-megagpu-16g 16

以下示例展示了如何提交使用 Compute Engine 机器类型(挂接了 GPU)的作业。

配有 GPU 的机器类型

或者,您可以选择配有 GPU 的旧版机器类型,而不使用 acceleratorConfig

  • 将规模层级设置为 CUSTOM
  • 根据您的任务所需的 GPU 数量和加速器类型,将您的主实例工作器和您想要加速的任何其他任务类型(工作器、参数服务器或评估程序)配置为使用以下支持 GPU 的机器类型之一:
    • standard_gpu:1 个 NVIDIA Tesla K80 GPU
    • complex_model_m_gpu:4 个 NVIDIA Tesla K80 GPU
    • complex_model_l_gpu:8 个 NVIDIA Tesla K80 GPU
    • standard_p100:1 个 NVIDIA Tesla P100 GPU
    • complex_model_m_p100:4 个 NVIDIA Tesla P100 GPU
    • standard_v100:1 个 NVIDIA Tesla V100 GPU
    • large_model_v100:1 个 NVIDIA Tesla V100 GPU
    • complex_model_m_v100:4 个 NVIDIA Tesla V100 GPU
    • complex_model_l_v100:8 个 NVIDIA Tesla V100 GPU

以下示例展示了如何通过 gcloud 命令使用支持 GPU 的机器类型提交作业。

如需了解详情,请参阅 AI Platform Training 的机器类型

支持 GPU 的区域

您必须在支持 GPU 的区域中运行作业。以下区域目前支持使用 GPU:

  • us-west1
  • us-west2
  • us-central1
  • us-east1
  • us-east4
  • northamerica-northeast1
  • southamerica-east1
  • europe-west1
  • europe-west2
  • europe-west4
  • asia-south1
  • asia-southeast1
  • asia-east1
  • asia-northeast1
  • asia-northeast3
  • australia-southeast1

此外,其中一些区域仅支持使用特定类型的 GPU。如需全面了解 AI Platform Training 服务(包括模型训练和在线/批量预测)的可用区域,请参阅区域指南

如果训练作业使用多种类型的 GPU,则它们必须位于您所在区域中的同一个地区。举例来说,如果您的主实例工作器使用 NVIDIA Tesla T4 GPU、参数服务器使用 NVIDIA Tesla K80 GPU,并且工作器使用 NVIDIA Tesla P100 GPU,则无法在 us-central1 中运行作业。虽然 us-central1 中提供所有这些 GPU 以用于训练作业,但该区域中没有一个地区能提供全部三种类型的 GPU。如需详细了解 GPU 的地区可用性,请参阅比较用于处理计算工作负载的 GPU

提交训练作业

您可以使用 gcloud ai-platform jobs submit training 命令提交训练作业。

  1. 定义一个描述所需 GPU 的 config.yaml 文件。YAML 文件的结构表示作业资源。以下是 config.yaml 文件的两个示例。

    第一个示例展示了使用 Compute Engine 机器类型(其中一些机器类型挂接了 GPU)的训练作业的配置文件:

    trainingInput:
      scaleTier: CUSTOM
      # Configure a master worker with 4 T4 GPUs
      masterType: n1-highcpu-16
      masterConfig:
        acceleratorConfig:
          count: 4
          type: NVIDIA_TESLA_T4
      # Configure 9 workers, each with 4 T4 GPUs
      workerCount: 9
      workerType: n1-highcpu-16
      workerConfig:
        acceleratorConfig:
          count: 4
          type: NVIDIA_TESLA_T4
      # Configure 3 parameter servers with no GPUs
      parameterServerCount: 3
      parameterServerType: n1-highmem-8
    

    下一个示例展示了配置与上述示例类似的作业的配置文件。但是,此配置使用的旧版机器类型配有 GPU 而不是使用 acceleratorConfig 挂接 GPU:

    trainingInput:
      scaleTier: CUSTOM
      # Configure a master worker with 4 K80 GPUs
      masterType: complex_model_m_gpu
      # Configure 9 workers, each with 4 K80 GPUs
      workerCount: 9
      workerType: complex_model_m_gpu
      # Configure 3 parameter servers with no GPUs
      parameterServerCount: 3
      parameterServerType: large_model
    
  2. 使用 gcloud 命令提交作业(需要在命令中添加指向 config.yaml 文件的 --config 参数)。以下示例假设您已为一些参数的值设置了环境变量(由 $ 符号后跟大写字母表示):

    gcloud ai-platform jobs submit training $JOB_NAME \
            --package-path $APP_PACKAGE_PATH \
            --module-name $MAIN_APP_MODULE \
            --job-dir $JOB_DIR \
            --region us-central1 \
            --config config.yaml \
            -- \
            --user_arg_1 value_1 \
             ...
            --user_arg_n value_n
    

或者,您也可以使用命令行标志(而不是在配置文件中)指定集群配置详细信息。详细了解如何使用这些标志

以下示例展示了不使用 config.yaml 文件如何提交配置与上一示例相同的作业(使用挂接了 GPU 的 Compute Engine 机器类型):

gcloud ai-platform jobs submit training $JOB_NAME \
        --package-path $APP_PACKAGE_PATH \
        --module-name $MAIN_APP_MODULE \
        --job-dir $JOB_DIR \
        --region us-central1 \
        --scale-tier custom \
        --master-machine-type n1-highcpu-16 \
        --master-accelerator count=4,type=nvidia-tesla-t4 \
        --worker-count 9 \
        --worker-machine-type n1-highcpu-16 \
        --worker-accelerator count=4,type=nvidia-tesla-t4 \
        --parameter-server-count 3 \
        --parameter-server-machine-type n1-highmem-8 \
        -- \
        --user_arg_1 value_1 \
         ...
        --user_arg_n value_n

注意:

  • 如果您在配置文件 (config.yaml) 中指定了一个选项,同时又将该选项指定为命令行标志,则命令行中的值将覆盖配置文件中的值。
  • -- 标志标记了 gcloud 特有标志的结束以及您要传递至应用的 USER_ARGS 的开始。
  • AI Platform Training 特有的标志(例如 --module-name--runtime-version--job-dir)必须位于空 -- 标志之前。AI Platform Training 服务可解读这些标志。
  • 如果指定 --job-dir 标志,则该标志必须位于空 -- 标志之前,原因是 AI Platform Training 会使用 --job-dir 来验证路径。
  • 如果指定 --job-dir 标志,您的应用也必须对该标志进行处理。虽然 --job-dir 位于空 -- 之前,该标志也会作为命令行标志传递至您的应用。
  • 您可以根据需要定义任意数量的 USER_ARGS。AI Platform Training 会将 --user_first_arg--user_second_arg 等传递至您的应用。

如需详细了解作业提交选项,请参阅启动训练作业指南。

将训练代码调整为使用 GPU

如果您在 TensorFlow 训练作业中使用 KerasEstimator,并希望使用具有一个 GPU 的单个虚拟机进行训练,则无需为 GPU 自定义代码。

如果您的训练集群包含多个 GPU,请在训练代码中使用 tf.distribute.Strategy API:

  • 对于具有多个 GPU 的单个虚拟机的训练,建议使用 TensorFlow 2.1 及更高版本中 Keras 完全支持的 MirroredStrategy
  • 要针对具有 GPU 的多个虚拟机进行训练,请参阅分布式训练的建议

要自定义 TensorFlow 向 GPU 分配特定操作的方式,请参阅使用 GPU 的 TensorFlow 指南。在这种情况下,您可能还需要参考了解 AI Platform Training 如何在每个虚拟机上设置 TF_CONFIG 环境变量

GPU 设备字符串

standard_gpu 机器的单个 GPU 标识为 "/gpu:0"。具有多个 GPU 的机器使用从 "/gpu:0" 开始的标识符,然后是 "/gpu:1",以此类推。例如,complex_model_m_gpu 机器具有四个 GPU,标识为 "/gpu:0""/gpu:3"

支持 GPU 的机器上的 Python 软件包

支持 GPU 的机器已预安装 tensorflow-gpu,这是提供 GPU 支持的 TensorFlow Python 软件包。如需查看所有预安装软件包的列表,请参阅运行时版本列表

维护事件

运行 AI Platform Training 作业的支持 GPU 的虚拟机偶尔需要进行 Compute Engine 主机维护。虚拟机配置为在此类维护事件结束后自动重启,但您可能需要执行一些额外工作才能确保作业能够灵活应对这些关停问题。对训练应用进行配置,以定期保存模型检查点(通常保存在您通过 gcloud ai-platform jobs submit training--job-dir 参数指定的 Cloud Storage 路径),并在某检查点已存在的情况下恢复最新的检查点。

您只需指定一个 model_dirTensorFlow Estimator 即可为您实现此功能。Estimator 定期将检查点保存到 model_dir 并尝试从最新的检查点进行加载,这样您就不必担心 GPU 工作器上的维护事件了。

如果您正在使用 Keras 进行训练,请使用 ModelCheckpoint 回调定期保存训练进度。如果您将 tf.distribute.Strategy 与 Keras 搭配使用,则虚拟机会使用检查点自动从重启中恢复。否则,将逻辑添加到您的训练代码中,以检查是否存在最近的检查点从存在的检查点恢复

如需了解更多高级案例,请参阅检查点的 TensorFlow 指南

后续步骤