使用 PyTorch 训练机器学习模型

本教程介绍如何运行使用 PyTorch 机器学习框架的训练作业。本教程介绍将作业配置为使用 PyTorch 与使用 AI Platform Training 支持的其他机器学习框架的细微差异。然后向您展示如何使用示例 PyTorch 代码运行训练作业,该示例代码根据芝加哥出租车行程数据集中的数据训练模型。

本教程还介绍如何将 PyTorch 与 GPU超参数调节搭配使用。

PyTorch 容器

AI Platform Training 的运行时版本不包含 PyTorch 作为依赖项。要运行使用 PyTorch 的训练作业,请指定一个预构建的 PyTorch 容器供 AI Platform Training 使用。

为训练配置预构建容器时,使用的某些语法与配置自定义容器的语法相同。但是,您不必创建自己的 Docker 容器;请改为指定 AI Platform 提供的容器映像的 URI,并提供您创建的 Python 训练软件包

AI Platform 提供以下预构建的 PyTorch 容器:

容器映像 URI PyTorch 版本 支持的处理器
gcr.io/cloud-ml-public/training/pytorch-xla.1-11 1.11 CPU、TPU
gcr.io/cloud-ml-public/training/pytorch-gpu.1-11 1.11 GPU
gcr.io/cloud-ml-public/training/pytorch-xla.1-10 1.10 CPU、TPU
gcr.io/cloud-ml-public/training/pytorch-gpu.1-10 1.10 GPU
gcr.io/cloud-ml-public/training/pytorch-xla.1-9 1.9 CPU、TPU
gcr.io/cloud-ml-public/training/pytorch-gpu.1-9 1.9 GPU
gcr.io/cloud-ml-public/training/pytorch-xla.1-7 1.7 CPU、TPU
gcr.io/cloud-ml-public/training/pytorch-gpu.1-7 1.7 GPU
gcr.io/cloud-ml-public/training/pytorch-xla.1-6 1.6 CPU、TPU
gcr.io/cloud-ml-public/training/pytorch-gpu.1-6 1.6 GPU
gcr.io/cloud-ml-public/training/pytorch-cpu.1-4 1.4 CPU
gcr.io/cloud-ml-public/training/pytorch-gpu.1-4 1.4 GPU

这些容器映像派生自 Deep Learning Containers,并包含 Deep Learning Containers 提供的依赖项

如果您要使用预构建容器中没有的 PyTorch 版本,请按照自定义容器使用指南进行操作。

准备工作

  1. 登录您的 Google Cloud 账号。如果您是 Google Cloud 新手,请创建一个账号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金,用于运行、测试和部署工作负载。
  2. 在 Google Cloud Console 中的项目选择器页面上,选择或创建一个 Google Cloud 项目

    转到“项目选择器”

  3. 确保您的 Google Cloud 项目已启用结算功能

  4. 启用 AI Platform Training & Prediction API。

    启用 API

  5. 安装 Google Cloud CLI。
  6. 如需初始化 gcloud CLI,请运行以下命令:

    gcloud init
  7. 在 Google Cloud Console 中的项目选择器页面上,选择或创建一个 Google Cloud 项目

    转到“项目选择器”

  8. 确保您的 Google Cloud 项目已启用结算功能

  9. 启用 AI Platform Training & Prediction API。

    启用 API

  10. 安装 Google Cloud CLI。
  11. 如需初始化 gcloud CLI,请运行以下命令:

    gcloud init

下载示例代码

运行以下命令下载示例 PyTorch 训练应用,然后导航至训练应用所在的目录:

git clone --depth=1 \
  https://github.com/GoogleCloudPlatform/ai-platform-samples.git

cd ai-platform-samples/training/pytorch/structured/python_package

(可选)检查训练代码的结构:

ls -pR

trainer/ 目录包含 PyTorch 训练应用,而 setup.py 提供了打包训练应用的配置详情。

创建 Cloud Storage 存储桶

创建 Cloud Storage 存储桶以存储封装的训练代码和训练作业创建的模型工件。运行以下命令:

gsutil mb -l us-central1 gs://BUCKET_NAME

BUCKET_NAME 替换为您为存储桶选择的唯一名称。了解存储桶名称要求

或者,您可以使用 Google Cloud 项目中的现有 Cloud Storage 存储桶。在本教程中,请务必使用 us-central1 区域中的存储桶。

训练 PyTorch 模型

本教程介绍了在 AI Platform Training 上训练 PyTorch 模型的几种方法:

  • 在具有 CPU 处理器的虚拟机实例上
  • 在具有 GPU 处理器的虚拟机上
  • 使用超参数调节(在具有 CPU 处理器的虚拟机上)

现在请选择上述方法之一,然后在本教程的余下部分中根据相应标签页中的说明操作。如果您想尝试使用其他配置进行训练,则可以重复本部分。

准备创建训练作业

在创建训练作业之前,请确保您的训练代码已准备就绪,并在本地环境中指定一些配置选项。

CPU

设置创建训练作业时要使用的多个 Bash 变量:

BUCKET_NAME=BUCKET_NAME
JOB_NAME=getting_started_pytorch_cpu
JOB_DIR=gs://${BUCKET_NAME}/${JOB_NAME}/models

BUCKET_NAME 替换为您在上一部分中创建的 Cloud Storage 存储桶的名称。

GPU

  1. 确保 PyTorch 训练代码知道训练作业使用的虚拟机上的 GPU,以便 PyTorch 将张量和模块适当地移动到 GPU。

    如果您使用提供的示例代码,则无需执行任何操作,因为示例代码包含检测运行代码的机器是否具有 GPU 的逻辑:

    cuda_availability = torch.cuda.is_available()
    if cuda_availability:
      device = torch.device('cuda:{}'.format(torch.cuda.current_device()))
    else:
      device = 'cpu'

    如果您更改了训练代码,请阅读 CUDA 语义的 PyTorch 指南,以确保使用 GPU。

  2. 设置创建训练作业时要使用的多个 Bash 变量:

    BUCKET_NAME=BUCKET_NAME
    JOB_NAME=getting_started_pytorch_gpu
    JOB_DIR=gs://${BUCKET_NAME}/${JOB_NAME}/models
    

    BUCKET_NAME 替换为您在上一部分中创建的 Cloud Storage 存储桶的名称。

超参数调节

本教程的示例代码调节学习速率和批次大小参数,以尽可能减少测试损失。

  1. 确保您的训练代码已准备好在 AI Platform Training 上使用超参数调节:

  2. 运行以下命令创建一个 config.yaml 文件,该文件指定超参数调节选项

    cat > config.yaml <<END
    trainingInput:
      hyperparameters:
        goal: MINIMIZE
        hyperparameterMetricTag: test_loss
        maxTrials: 2
        maxParallelTrials: 2
        enableTrialEarlyStopping: True
        params:
        - parameterName: learning-rate
          type: DOUBLE
          minValue: 0.0001
          maxValue: 1
          scaleType: UNIT_LOG_SCALE
        - parameterName: batch-size
          type: INTEGER
          minValue: 1
          maxValue: 256
          scaleType: UNIT_LINEAR_SCALE
    END
    

    这些选项调节 --learning-rate--batch-size 超参数,以尽可能减少模型损失。

  3. 设置创建训练作业时要使用的多个 Bash 变量:

    BUCKET_NAME=BUCKET_NAME
    JOB_NAME=getting_started_pytorch_hptuning
    JOB_DIR=gs://${BUCKET_NAME}/${JOB_NAME}/models
    

    BUCKET_NAME 替换为您在上一部分中创建的 Cloud Storage 存储桶的名称。

这些 Bash 变量具有以下用途:

  • JOB_NAME 是您的 AI Platform Training 作业的标识符。它在您的 Google Cloud 项目中的 AI Platform Training 作业中必须具有唯一性。
  • AI Platform Training 使用 JOB_DIR 来确定上传训练应用的准确位置。完成训练后,训练应用还会使用 JOB_DIR 确定导出模型工件的位置。

创建训练作业

运行以下命令以创建训练作业:

CPU

gcloud ai-platform jobs submit training ${JOB_NAME} \
  --region=us-central1 \
  --master-image-uri=gcr.io/cloud-ml-public/training/pytorch-xla.1-10 \
  --scale-tier=BASIC \
  --job-dir=${JOB_DIR} \
  --package-path=./trainer \
  --module-name=trainer.task \
  -- \
  --train-files=gs://cloud-samples-data/ai-platform/chicago_taxi/training/small/taxi_trips_train.csv \
  --eval-files=gs://cloud-samples-data/ai-platform/chicago_taxi/training/small/taxi_trips_eval.csv \
  --num-epochs=10 \
  --batch-size=100 \
  --learning-rate=0.001

GPU

gcloud ai-platform jobs submit training ${JOB_NAME} \
  --region=us-central1 \
  --master-image-uri=gcr.io/cloud-ml-public/training/pytorch-gpu.1-10 \
  --scale-tier=CUSTOM \
  --master-machine-type=n1-standard-8 \
  --master-accelerator=type=nvidia-tesla-p100,count=1 \
  --job-dir=${JOB_DIR} \
  --package-path=./trainer \
  --module-name=trainer.task \
  -- \
  --train-files=gs://cloud-samples-data/ai-platform/chicago_taxi/training/small/taxi_trips_train.csv \
  --eval-files=gs://cloud-samples-data/ai-platform/chicago_taxi/training/small/taxi_trips_eval.csv \
  --num-epochs=10 \
  --batch-size=100 \
  --learning-rate=0.001

超参数调节

gcloud ai-platform jobs submit training ${JOB_NAME} \
  --region=us-central1 \
  --master-image-uri=gcr.io/cloud-ml-public/training/pytorch-xla.1-10 \
  --scale-tier=BASIC \
  --job-dir=${JOB_DIR} \
  --package-path=./trainer \
  --module-name=trainer.task \
  --config=config.yaml \
  -- \
  --train-files=gs://cloud-samples-data/ai-platform/chicago_taxi/training/small/taxi_trips_train.csv \
  --eval-files=gs://cloud-samples-data/ai-platform/chicago_taxi/training/small/taxi_trips_eval.csv \
  --num-epochs=10

阅读训练作业指南,了解配置标志以及如何使用它们来自定义训练。

该命令会返回类似于如下所示的消息:

Job [JOB_NAME] submitted successfully.
Your job is still active. You may view the status of your job with the command

  $ gcloud ai-platform jobs describe JOB_NAME

or continue streaming the logs with the command

  $ gcloud ai-platform jobs stream-logs JOB_NAME
jobId: JOB_NAME
state: QUEUED

您可以使用以下命令监控作业状态:

gcloud ai-platform jobs describe ${JOB_NAME}

您可以使用以下命令流式传输作业的训练日志:

gcloud ai-platform jobs stream-logs ${JOB_NAME}

训练作业完成后,它会将经过训练的机器学习模型保存到名为 model.pth 文件中,该文件位于您指定的 JOB_DIR Cloud Storage 目录中的一个带时间戳的目录下。

后续步骤