使用 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. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. Enable the AI Platform Training & Prediction API.

    Enable the API

  5. Install the Google Cloud CLI.
  6. To initialize the gcloud CLI, run the following command:

    gcloud init
  7. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  8. Make sure that billing is enabled for your Google Cloud project.

  9. Enable the AI Platform Training & Prediction API.

    Enable the API

  10. Install the Google Cloud CLI.
  11. To initialize the gcloud CLI, run the following command:

    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 存储分区以存储封装的训练代码和训练作业创建的模型工件。运行以下命令:

gcloud storage buckets create gs://BUCKET_NAME --location=us-central1

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 目录中的一个带时间戳的目录下。

后续步骤