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

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. 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, 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. Make sure that billing is enabled for your Google Cloud project.

  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
    

后续步骤