使用自定义容器训练机器学习模型

AI Platform Training 支持在自定义容器中进行训练,可让用户在 AI Platform Training 上运行预安装了任何机器学习框架或算法的自定义 Docker 容器。本教程提供有关如何在 AI Platform Training 上使用自定义容器训练 PyTorch 模型的入门级演示。

概览

本入门指南演示了在 AI Platform Training 上使用自定义容器进行训练的过程。此训练过程使用了一个基于 MNIST 数据集对手写数字进行分类的基本模型。

本指南介绍以下步骤:

  • 设置项目和本地环境
  • 创建自定义容器
    • 编写 Dockerfile
    • 在本地构建并测试 Docker 映像
  • 将映像推送到 Container Registry
  • 提交自定义容器训练作业
  • 提交超参数调节作业
  • 将 GPU 与自定义容器搭配使用

准备工作

对于此入门指南,请使用安装了 Google Cloud CLI 的任何环境。

可选:查看有关使用自定义容器进行训练的概念信息

完成以下步骤以设置 GCP 账号,启用所需的 API,以及安装和激活 Cloud SDK。

  1. 登录您的 Google Cloud 账号。如果您是 Google Cloud 新手,请创建一个账号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金,用于运行、测试和部署工作负载。
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

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

  4. Enable the AI Platform Training & Prediction, Compute Engine and Container Registry APIs.

    Enable the APIs

  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. 确保您的 Google Cloud 项目已启用结算功能

  9. Enable the AI Platform Training & Prediction, Compute Engine and Container Registry APIs.

    Enable the APIs

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

    gcloud init
  12. 安装 Docker

    如果您使用的是基于 Linux 的操作系统(例如 Ubuntu 或 Debian),请将您的用户名添加到 docker 群组,这样您就可以在不使用 sudo 的情况下运行 Docker:

    sudo usermod -a -G docker ${USER}

    在将您自己添加到 docker 群组之后,您可能需要重启系统。

  13. 打开 Docker。如需确保 Docker 正在运行,请运行以下 Docker 命令以返回当前时间和日期:
    docker run busybox date
  14. 使用 gcloud 作为 Docker 的凭据帮助程序:
    gcloud auth configure-docker
  15. 可选:如果要在本地使用 GPU 运行容器,请安装 nvidia-docker

设置 Cloud Storage 存储分区

本部分介绍如何创建新存储桶。您可以使用现有存储桶,但它所在区域必须与您计划运行 AI Platform 作业的区域相同。此外,如果该存储桶不属于您用于运行 AI Platform Training 的项目,则您必须明确向 AI Platform Training 服务账号授予访问权限

  1. 为新存储桶指定名称。该名称在 Cloud Storage 的所有存储桶中必须是唯一的。

    BUCKET_NAME="YOUR_BUCKET_NAME"

    例如,使用附加了 -aiplatform 的项目名称:

    PROJECT_ID=$(gcloud config list project --format "value(core.project)")
    BUCKET_NAME=${PROJECT_ID}-aiplatform
  2. 检查您创建的存储桶名称。

    echo $BUCKET_NAME
  3. 为您的存储桶选择一个区域,并设置 REGION 环境变量。

    使用您计划在其中运行 AI Platform Training 作业的区域。请参阅 AI Platform Training 服务的可用区域

    例如,以下代码会创建 REGION 并将其设置为 us-central1

    REGION=us-central1
  4. 创建新的存储分区:

    gcloud storage buckets create gs://$BUCKET_NAME --location=$REGION

下载本教程的代码

  1. 输入以下命令来下载 AI Platform Training 示例 ZIP 文件:

    wget https://github.com/GoogleCloudPlatform/cloudml-samples/archive/master.zip
    
  2. 解压缩文件,以提取 cloudml-samples-master 目录。

    unzip master.zip
    
  3. 导航到 cloudml-samples-master > pytorch > containers > quickstart > mnist 目录。本演示中的命令必须从 mnist 目录运行。

    cd cloudml-samples-master/pytorch/containers/quickstart/mnist
    

创建自定义容器

要创建自定义容器,首先要定义 Dockerfile 以安装训练作业所需的依赖项。接着,在本地构建和测试 Docker 映像以对其进行验证。最后将该映像与 AI Platform Training 一起使用。

编写 Dockerfile

本教程中提供的示例 Dockerfile 完成以下步骤:

  1. 使用具有内置 Python 依赖项的 Python 2.7 基础映像。
  2. 安装其他依赖项,包括 PyTorch、gcloud CLI 和用于超参数调节的 cloudml-hypertune
  3. 将训练应用的代码复制到容器中。
  4. 配置 AI Platform Training 的入口点,以在容器启动时运行训练代码。

您的 Dockerfile 可能包含其他逻辑,具体取决于您的需求。详细了解如何编写 Dockerfile

# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the \"License\");
# you may not use this file except in compliance with the License.\n",
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an \"AS IS\" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Dockerfile
FROM python:2.7.16-jessie
WORKDIR /root

# Installs pytorch and torchvision.
RUN pip install torch==1.0.0 torchvision==0.2.1

# Installs cloudml-hypertune for hyperparameter tuning.
# It’s not needed if you don’t want to do hyperparameter tuning.
RUN pip install cloudml-hypertune

# Installs google cloud sdk, this is mostly for using gsutil to export model.
RUN wget -nv \
    https://dl.google.com/dl/cloudsdk/release/google-cloud-sdk.tar.gz && \
    mkdir /root/tools && \
    tar xvzf google-cloud-sdk.tar.gz -C /root/tools && \
    rm google-cloud-sdk.tar.gz && \
    /root/tools/google-cloud-sdk/install.sh --usage-reporting=false \
        --path-update=false --bash-completion=false \
        --disable-installation-options && \
    rm -rf /root/.config/* && \
    ln -s /root/.config /config && \
    # Remove the backup directory that gcloud creates
    rm -rf /root/tools/google-cloud-sdk/.install/.backup

# Path configuration
ENV PATH $PATH:/root/tools/google-cloud-sdk/bin
# Make sure gsutil will use the default service account
RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg

# Copies the trainer code 
RUN mkdir /root/trainer
COPY trainer/mnist.py /root/trainer/mnist.py

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "trainer/mnist.py"]

在本地构建并测试 Docker 映像

  1. 使用环境变量创建正确的映像 URI,并构建 Docker 映像。-t 标志使用您选择的 IMAGE_REPO_NAMEIMAGE_TAG 命名和标记映像。您可以为映像选择不同的名称和标记。

    export PROJECT_ID=$(gcloud config list project --format "value(core.project)")
    export IMAGE_REPO_NAME=mnist_pytorch_custom_container
    export IMAGE_TAG=mnist_pytorch_cpu
    export IMAGE_URI=gcr.io/$PROJECT_ID/$IMAGE_REPO_NAME:$IMAGE_TAG
    
    docker build -f Dockerfile -t $IMAGE_URI ./
  2. 通过在新的容器中本地运行映像进行验证。请注意,--epochs 标志将传递给训练程序脚本。

    docker run $IMAGE_URI --epochs 1

将映像推送到 Container Registry

如果本地运行有效,您可以将 Docker 映像推送到项目的 Container Registry。

在此之前,请先运行 gcloud auth configure-docker(如果您尚未这样做)。

docker push $IMAGE_URI

提交并监控作业

  1. 为您的作业请求定义环境变量。

    • MODEL_DIR 会在您的 Cloud Storage 存储分区中命名一个新的带有时间戳的目录,训练完成后,您保存的模型文件将存储在该目录中。
    • REGION 指定 AI Platform Training 训练的有效区域
    export MODEL_DIR=pytorch_model_$(date +%Y%m%d_%H%M%S)
    export REGION=us-central1
    export JOB_NAME=custom_container_job_$(date +%Y%m%d_%H%M%S)
    
  2. 使用 gcloud CLI 将训练作业提交到 AI Platform Training。 使用 --master-image-uri 标志将 URI 传递至 Docker 映像:

    gcloud ai-platform jobs submit training $JOB_NAME \
      --region $REGION \
      --master-image-uri $IMAGE_URI \
      -- \
      --model-dir=gs://$BUCKET_NAME/$MODEL_DIR \
      --epochs=10
    
  3. 提交作业后,您可以监控作业状态和流日志:

    gcloud ai-platform jobs describe $JOB_NAME
    gcloud ai-platform jobs stream-logs $JOB_NAME
    

提交超参数调节作业

还需要对超参数调节作业进行一些调整。请注意示例代码中的这些部分:

  • 示例 Dockerfile 包含 cloudml-hypertune 软件包,以便将其安装在自定义容器中。
  • 示例代码 (mnist.py):
    • 使用 cloudml-hypertune,以通过调用其辅助函数 report_hyperparameter_tuning_metric 来报告每次试验的结果。除非作业未作为超参数调节作业提交,否则示例代码会在评估后报告超参数调节结果。
    • 为每个超参数添加命令行参数,并使用 argparse 处理参数解析。
  • 作业请求包括 TrainingInput 对象中的 HyperparameterSpec。在这种情况下,我们会调整 --lr--momentum,以尽可能减少模型损失。
  1. 创建 config.yaml 文件以定义超参数规范。重新定义 MODEL_DIRJOB_NAME。请定义 REGION(如果您尚未这样做):

    export MODEL_DIR=pytorch_hptuning_model_$(date +%Y%m%d_%H%M%S)
    export REGION=us-central1
    export JOB_NAME=custom_container_job_hptuning_$(date +%Y%m%d_%H%M%S)
    
    # Creates a YAML file with job request.
    cat > config.yaml <<EOF
    trainingInput:
      hyperparameters:
        goal: MINIMIZE
        hyperparameterMetricTag: "my_loss"
        maxTrials: 20
        maxParallelTrials: 5
        enableTrialEarlyStopping: True
        params:
        - parameterName: lr
          type: DOUBLE
          minValue: 0.0001
          maxValue: 0.1
        - parameterName: momentum
          type: DOUBLE
          minValue: 0.2
          maxValue: 0.8
    EOF
    
  2. 将超参数调节作业提交至 AI Platform Training:

    gcloud ai-platform jobs submit training $JOB_NAME \
      --scale-tier BASIC \
      --region $REGION \
      --master-image-uri $IMAGE_URI \
      --config config.yaml \
      -- \
      --epochs=5 \
      --model-dir="gs://$BUCKET_NAME/$MODEL_DIR"
    

将 GPU 与自定义容器搭配使用

要使用 GPU 提交自定义容器作业,您必须构建与之前使用的 Docker 映像不同的Docker 映像。我们提供了一个与 GPU 搭配使用的 Dockerfile 示例,它满足以下要求:

  • 在容器中预安装 CUDA 工具包和 cuDNN。推荐使用 nvidia/cuda 映像作为基础映像来执行此操作,因为它预安装了 CUDA 工具包和 cuDNN,并且有助于您正确设置相关的环境变量。
  • 安装其他依赖项,比如 wgetcurlpip 以及训练应用所需的任何其他依赖项。
# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the \"License\");
# you may not use this file except in compliance with the License.\n",
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an \"AS IS\" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Dockerfile-gpu
FROM nvidia/cuda:9.0-cudnn7-runtime

# Installs necessary dependencies.
RUN apt-get update && apt-get install -y --no-install-recommends \
         wget \
         curl \
         python-dev && \
     rm -rf /var/lib/apt/lists/*

# Installs pip.
RUN curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && \
    python get-pip.py && \
    pip install setuptools && \
    rm get-pip.py

WORKDIR /root

# Installs pytorch and torchvision.
RUN pip install torch==1.0.0 torchvision==0.2.1

# Installs cloudml-hypertune for hyperparameter tuning.
# It’s not needed if you don’t want to do hyperparameter tuning.
RUN pip install cloudml-hypertune

# Installs google cloud sdk, this is mostly for using gsutil to export model.
RUN wget -nv \
    https://dl.google.com/dl/cloudsdk/release/google-cloud-sdk.tar.gz && \
    mkdir /root/tools && \
    tar xvzf google-cloud-sdk.tar.gz -C /root/tools && \
    rm google-cloud-sdk.tar.gz && \
    /root/tools/google-cloud-sdk/install.sh --usage-reporting=false \
        --path-update=false --bash-completion=false \
        --disable-installation-options && \
    rm -rf /root/.config/* && \
    ln -s /root/.config /config && \
    # Remove the backup directory that gcloud creates
    rm -rf /root/tools/google-cloud-sdk/.install/.backup

# Path configuration
ENV PATH $PATH:/root/tools/google-cloud-sdk/bin
# Make sure gsutil will use the default service account
RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg

# Copies the trainer code 
RUN mkdir /root/trainer
COPY trainer/mnist.py /root/trainer/mnist.py

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "trainer/mnist.py"]

在本地构建并测试 GPU Docker 映像

  1. 使用 GPU Dockerfile 为 GPU 训练作业构建新映像。为避免替换 CPU 映像,必须使用与本教程前面使用的名称不同的名称重新定义 IMAGE_REPO_NAMEIMAGE_TAG

    export PROJECT_ID=$(gcloud config list project --format "value(core.project)")
    export IMAGE_REPO_NAME=mnist_pytorch_gpu_container
    export IMAGE_TAG=mnist_pytorch_gpu
    export IMAGE_URI=gcr.io/$PROJECT_ID/$IMAGE_REPO_NAME:$IMAGE_TAG
    
    docker build -f Dockerfile-gpu -t $IMAGE_URI ./
  2. 如果您的机器上有 GPU,并且已安装 nvidia-docker,则可以通过在本地运行映像来进行验证:

    docker run --runtime=nvidia $IMAGE_URI --epochs 1
  3. 将 Docker 映像推送到 Container Registry。在此之前,请先运行 gcloud auth configure-docker(如果您尚未这样做)。

    docker push $IMAGE_URI

提交作业

此示例使用基本 GPU 调节层级来提交训练作业请求。请参阅其他机器选项以了解使用 GPU 进行训练。

  1. 重新定义 MODEL_DIRJOB_NAME。请定义 REGION(如果您尚未这样做):

    export MODEL_DIR=pytorch_model_gpu_$(date +%Y%m%d_%H%M%S)
    export REGION=us-central1
    export JOB_NAME=custom_container_job_gpu_$(date +%Y%m%d_%H%M%S)
    
  2. 使用 gcloud CLI 将训练作业提交到 AI Platform Training。 使用 --master-image-uri 标志将 URI 传递至 Docker 映像。

    gcloud ai-platform jobs submit training $JOB_NAME \
      --scale-tier BASIC_GPU \
      --region $REGION \
      --master-image-uri $IMAGE_URI \
      -- \
      --epochs=5 \
      --model-dir=gs://$BUCKET_NAME/$MODEL_DIR
    

后续步骤