使用 Azure 机器学习大规模训练 PyTorch 模型Train PyTorch models at scale with Azure Machine Learning

本文介绍了如何使用 Azure 机器学习在企业范围内运行 PyTorch 训练脚本。In this article, learn how to run your PyTorch training scripts at enterprise scale using Azure Machine Learning.

本文中的示例脚本用来对鸡和火鸡图像进行分类,以基于 PyTorch 的迁移学习教程构建深度学习神经网络 (DNN)。The example scripts in this article are used to classify chicken and turkey images to build a deep learning neural network (DNN) based on PyTorch's transfer learning tutorial.

无论是从头开始训练深度学习 PyTorch 模型,还是将现有模型引入云中,都可以通过 Azure 机器学习使用弹性云计算资源来横向扩展开源训练作业。Whether you're training a deep learning PyTorch model from the ground-up or you're bringing an existing model into the cloud, you can use Azure Machine Learning to scale out open-source training jobs using elastic cloud compute resources. 你可以通过 Azure 机器学习来构建、部署和监视生产级模型以及对其进行版本控制。You can build, deploy, version, and monitor production-grade models with Azure Machine Learning.

先决条件Prerequisites

在以下任一环境中运行此代码:Run this code on either of these environments:

  • Azure 机器学习计算实例 - 无需下载或安装Azure Machine Learning compute instance - no downloads or installation necessary

    • 在开始本教程之前完成教程:设置环境和工作区创建预先加载了 SDK 和示例存储库的专用笔记本服务器。Complete the Tutorial: Setup environment and workspace to create a dedicated notebook server pre-loaded with the SDK and the sample repository.
    • 在笔记本服务器上的示例深度学习文件夹中,通过导航到以下目录,找到已完成且已展开的笔记本:how-to-use-azureml > ml-frameworks > pytorch > train-hyperparameter-tune-deploy-with-pytorch 文件夹。In the samples deep learning folder on the notebook server, find a completed and expanded notebook by navigating to this directory: how-to-use-azureml > ml-frameworks > pytorch > train-hyperparameter-tune-deploy-with-pytorch folder.
  • 你自己的 Jupyter 笔记本服务器Your own Jupyter Notebook server

    此外还可以在 GitHub 示例页上查找本指南的完整 Jupyter Notebook 版本You can also find a completed Jupyter Notebook version of this guide on the GitHub samples page. 该笔记本包含扩展部分,其中涵盖智能超参数优化、模型部署和笔记本小组件。The notebook includes expanded sections covering intelligent hyperparameter tuning, model deployment, and notebook widgets.

设置试验Set up the experiment

本部分通过加载所需的 Python 包、初始化工作区、创建计算目标和定义训练环境来设置训练实验。This section sets up the training experiment by loading the required Python packages, initializing a workspace, creating the compute target, and defining the training environment.

导入程序包Import packages

首先,导入必需的 Python 库。First, import the necessary Python libraries.

import os
import shutil

from azureml.core.workspace import Workspace
from azureml.core import Experiment
from azureml.core import Environment

from azureml.core.compute import ComputeTarget, AmlCompute
from azureml.core.compute_target import ComputeTargetException

初始化工作区Initialize a workspace

Azure 机器学习工作区是服务的顶级资源。The Azure Machine Learning workspace is the top-level resource for the service. 它提供了一个集中的位置来处理创建的所有项目。It provides you with a centralized place to work with all the artifacts you create. 在 Python SDK 中,可以通过创建 workspace 对象来访问工作区项目。In the Python SDK, you can access the workspace artifacts by creating a workspace object.

根据在先决条件部分中创建的 config.json 文件创建工作区对象。Create a workspace object from the config.json file created in the prerequisites section.

ws = Workspace.from_config()

获取数据Get the data

该数据集针对火鸡和鸡分别包含了大约 120 个训练图像,每个类有 100 个验证图像。The dataset consists of about 120 training images each for turkeys and chickens, with 100 validation images for each class. 在训练脚本 pytorch_train.py 中,我们将下载并提取该数据集。We will download and extract the dataset as part of our training script pytorch_train.py. 这些图像是 Open Images v5 Dataset 的子集。The images are a subset of the Open Images v5 Dataset.

准备训练脚本Prepare training script

在本教程中,已提供了训练脚本 pytorch_train.pyIn this tutorial, the training script, pytorch_train.py, is already provided. 实际上,你可以原样接受任何自定义的训练脚本,并使用 Azure 机器学习运行它。In practice, you can take any custom training script, as is, and run it with Azure Machine Learning.

为训练脚本创建一个文件夹。Create a folder for your training script(s).

project_folder = './pytorch-birds'
os.makedirs(project_folder, exist_ok=True)
shutil.copy('pytorch_train.py', project_folder)

创建计算目标Create a compute target

创建用于运行 PyTorch 作业的计算目标。Create a compute target for your PyTorch job to run on. 在此示例中,创建启用了 GPU 的 Azure 机器学习计算群集。In this example, create a GPU-enabled Azure Machine Learning compute cluster.

cluster_name = "gpu-cluster"

try:
    compute_target = ComputeTarget(workspace=ws, name=cluster_name)
    print('Found existing compute target')
except ComputeTargetException:
    print('Creating a new compute target...')
    compute_config = AmlCompute.provisioning_configuration(vm_size='STANDARD_NC6', 
                                                           max_nodes=4)

    compute_target = ComputeTarget.create(ws, cluster_name, compute_config)

    compute_target.wait_for_completion(show_output=True, min_node_count=None, timeout_in_minutes=20)

有关计算目标的详细信息,请参阅什么是计算目标一文。For more information on compute targets, see the what is a compute target article.

定义环境Define your environment

若要定义封装训练脚本依赖项的 Azure ML 环境,可以定义自定义环境或使用和 Azure ML 特选环境。To define the Azure ML Environment that encapsulates your training script's dependencies, you can either define a custom environment or use an Azure ML curated environment.

使用特选环境Use a curated environment

如果不想定义自己的环境,Azure ML 会提供预生成的特选环境。Azure ML provides prebuilt, curated environments if you don't want to define your own environment. Azure ML 具有几个适用于 PyTorch 的 CPU 和 GPU 特选环境,这些环境对应不同版本的 PyTorch。Azure ML has several CPU and GPU curated environments for PyTorch corresponding to different versions of PyTorch. 有关详细信息,请参阅此文For more info, see here.

若要使用特选环境,可以改为运行以下命令:If you want to use a curated environment, you can run the following command instead:

curated_env_name = 'AzureML-PyTorch-1.6-GPU'
pytorch_env = Environment.get(workspace=ws, name=curated_env_name)

若要查看特选环境中包含的包,可以将 conda 依赖项写入磁盘:To see the packages included in the curated environment, you can write out the conda dependencies to disk:

pytorch_env.save_to_directory(path=curated_env_name)

确保特选环境包括训练脚本所需的所有依赖项。Make sure the curated environment includes all the dependencies required by your training script. 如果没有,则必须修改环境以包含缺少的依赖项。If not, you will have to modify the environment to include the missing dependencies. 请注意,如果修改了环境,则必须为它提供新名称,因为“AzureML”前缀是为特选环境保留的。Note that if the environment is modified, you will have to give it a new name, as the 'AzureML' prefix is reserved for curated environments. 如果修改了 conda 依赖项 YAML 文件,则可以使用新名称从该文件创建新环境,例如:If you modified the conda dependencies YAML file, you can create a new environment from it with a new name, e.g.:

pytorch_env = Environment.from_conda_specification(name='pytorch-1.6-gpu', file_path='./conda_dependencies.yml')

如果改为直接修改了特选环境对象,则可以使用新名称克隆该环境:If you had instead modified the curated environment object directly, you can clone that environment with a new name:

pytorch_env = pytorch_env.clone(new_name='pytorch-1.6-gpu')

创建自定义环境Create a custom environment

还可以创建自己的 Azure ML 环境,以封装训练脚本的依赖项。You can also create your own Azure ML environment that encapsulates your training script's dependencies.

首先,在 YAML 文件中定义 conda 依赖项;在本例中,该文件名为 conda_dependencies.ymlFirst, define your conda dependencies in a YAML file; in this example the file is named conda_dependencies.yml.

channels:
- conda-forge
dependencies:
- python=3.6.2
- pip:
  - azureml-defaults
  - torch==1.6.0
  - torchvision==0.7.0
  - future==0.17.1
  - pillow

基于此 conda 环境规范创建 Azure ML 环境。Create an Azure ML environment from this conda environment specification. 此环境将在运行时打包到 Docker 容器中。The environment will be packaged into a Docker container at runtime.

在默认情况下,如果未指定基础映像,Azure ML 将使用 CPU 映像 azureml.core.environment.DEFAULT_CPU_IMAGE 作为基础映像。By default if no base image is specified, Azure ML will use a CPU image azureml.core.environment.DEFAULT_CPU_IMAGE as the base image. 由于本示例在 GPU 群集上运行训练,因此你需要指定具有必要 GPU 驱动程序和依赖项的 GPU 基础映像。Since this example runs training on a GPU cluster, you will need to specify a GPU base image that has the necessary GPU drivers and dependencies. Azure ML 维护一组在 Microsoft 容器注册表 (MCR) 上发布的基础映像,你可以使用这些映像,请参阅 Azure/AzureML 容器 GitHub 存储库获取详细信息。Azure ML maintains a set of base images published on Microsoft Container Registry (MCR) that you can use, see the Azure/AzureML-Containers GitHub repo for more information.

pytorch_env = Environment.from_conda_specification(name='pytorch-1.6-gpu', file_path='./conda_dependencies.yml')

# Specify a GPU base image
pytorch_env.docker.enabled = True
pytorch_env.docker.base_image = 'mcr.microsoft.com/azureml/openmpi3.1.2-cuda10.1-cudnn7-ubuntu18.04'

提示

或者,也可以直接在自定义 Docker 映像或 Dockerfile 中捕获所有依赖项,然后从中创建环境。Optionally, you can just capture all your dependencies directly in a custom Docker image or Dockerfile, and create your environment from that. 有关详细信息,请参阅通过自定义映像进行训练For more information, see Train with custom image.

有关创建和使用环境的详细信息,请参阅在 Azure 机器学习中创建和使用软件环境For more information on creating and using environments, see Create and use software environments in Azure Machine Learning.

配置和提交训练运行Configure and submit your training run

创建 ScriptRunConfigCreate a ScriptRunConfig

创建一个 ScriptRunConfig 对象,以指定训练作业的配置详细信息,包括训练脚本、要使用的环境,以及要在其上运行的计算目标。Create a ScriptRunConfig object to specify the configuration details of your training job, including your training script, environment to use, and the compute target to run on. 如果在 arguments 参数中指定,训练脚本的任何参数都将通过命令行传递。Any arguments to your training script will be passed via command line if specified in the arguments parameter.

from azureml.core import ScriptRunConfig

src = ScriptRunConfig(source_directory=project_folder,
                      script='pytorch_train.py',
                      arguments=['--num_epochs', 30, '--output_dir', './outputs'],
                      compute_target=compute_target,
                      environment=pytorch_env)

警告

Azure 机器学习通过复制整个源目录来运行训练脚本。Azure Machine Learning runs training scripts by copying the entire source directory. 如果你有不想上传的敏感数据,请使用 .ignore 文件或不将其包含在源目录中。If you have sensitive data that you don't want to upload, use a .ignore file or don't include it in the source directory . 改为使用 Azure ML 数据集来访问数据。Instead, access your data using an Azure ML dataset.

有关通过 ScriptRunConfig 配置作业的详细信息,请参阅配置并提交训练运行For more information on configuring jobs with ScriptRunConfig, see Configure and submit training runs.

警告

如果你以前使用 PyTorch 估算器来配置 PyTorch 训练作业,请注意,在 Azure ML SDK 的未来版本中将会弃用估算器。If you were previously using the PyTorch estimator to configure your PyTorch training jobs, please note that Estimators will be deprecated in a future release of the Azure ML SDK. 对于不低于 1.15.0 版本的 Azure ML SDK,建议使用 ScriptRunConfig 作为配置训练作业(包括使用 DL 框架的作业)的方法。With Azure ML SDK >= 1.15.0, ScriptRunConfig is the recommended way to configure training jobs, including those using DL frameworks.

提交运行Submit your run

运行对象在作业运行时和运行后提供运行历史记录的接口。The Run object provides the interface to the run history while the job is running and after it has completed.

run = Experiment(ws, name='pytorch-birds').submit(src)
run.wait_for_completion(show_output=True)

在运行执行过程中发生的情况What happens during run execution

执行运行时,会经历以下阶段:As the run is executed, it goes through the following stages:

  • 准备:根据所定义的环境创建 docker 映像。Preparing: A docker image is created according to the environment defined. 将映像上传到工作区的容器注册表,缓存以用于后续运行。The image is uploaded to the workspace's container registry and cached for later runs. 还会将日志流式传输到运行历史记录,可以查看日志以监视进度。Logs are also streamed to the run history and can be viewed to monitor progress. 如果改为指定特选环境,将会使用支持该特选环境的缓存映像。If a curated environment is specified instead, the cached image backing that curated environment will be used.

  • 缩放:如果 Batch AI 群集执行运行所需的节点多于当前可用节点,则群集将尝试纵向扩展。Scaling: The cluster attempts to scale up if the Batch AI cluster requires more nodes to execute the run than are currently available.

  • 正在运行:将脚本文件夹中的所有脚本上传到计算目标,装载或复制数据存储,然后执行 scriptRunning: All scripts in the script folder are uploaded to the compute target, data stores are mounted or copied, and the script is executed. stdout 和 ./logs 文件夹中的输出会流式传输到运行历史记录,并可用于监视运行。Outputs from stdout and the ./logs folder are streamed to the run history and can be used to monitor the run.

  • 后期处理:将运行的 ./outputs 文件夹复制到运行历史记录。Post-Processing: The ./outputs folder of the run is copied over to the run history.

注册或下载模型Register or download a model

训练模型后,可以将其注册到工作区。Once you've trained the model, you can register it to your workspace. 凭借模型注册,可以在工作区中存储模型并对其进行版本控制,从而简化模型管理和部署Model registration lets you store and version your models in your workspace to simplify model management and deployment.

model = run.register_model(model_name='pytorch-birds', model_path='outputs/model.pt')

提示

部署指南包含有关模型注册的部分,但由于你已有一个已注册的模型,因而可以直接跳到创建计算目标进行部署。The deployment how-to contains a section on registering models, but you can skip directly to creating a compute target for deployment, since you already have a registered model.

此外,还可以使用“运行”对象下载模型的本地副本。You can also download a local copy of the model by using the Run object. 在训练脚本 pytorch_train.py 中,一个 PyTorch“保存”对象将模型保存到本地文件夹(计算目标的本地)。In the training script pytorch_train.py, a PyTorch save object persists the model to a local folder (local to the compute target). 可以使用“运行”对象下载副本。You can use the Run object to download a copy.

# Create a model folder in the current directory
os.makedirs('./model', exist_ok=True)

# Download the model from run history
run.download_file(name='outputs/model.pt', output_file_path='./model/model.pt'), 

分布式训练Distributed training

Azure 机器学习还支持多节点分布式 PyTorch 作业,以便可以缩放训练工作负荷。Azure Machine Learning also supports multi-node distributed PyTorch jobs so that you can scale your training workloads. 你可以轻松运行分布式 PyTorch 作业,Azure ML 将为你管理业务流程。You can easily run distributed PyTorch jobs and Azure ML will manage the orchestration for you.

Azure ML 支持使用 Horovod 和 PyTorch 的内置 DistributedDataParallel 模块来运行分布式 PyTorch 作业。Azure ML supports running distributed PyTorch jobs with both Horovod and PyTorch's built-in DistributedDataParallel module.

HorovodHorovod

Horovod 是 Uber 开发的用于分布式训练的开放源代码 all reduce 框架。Horovod is an open-source, all reduce framework for distributed training developed by Uber. 它提供了一个简单的路径来编写分布式 PyTorch 代码进行训练。It offers an easy path to writing distributed PyTorch code for training.

训练代码必须使用 Horovod 检测,以进行分布式训练。Your training code will have to be instrumented with Horovod for distributed training. 有关将 Horovod 与 PyTorch 配合使用的详细信息,请参阅 Horovod 文档For more information using Horovod with PyTorch, see the Horovod documentation.

此外,请确保训练环境包含 horovod 包。Additionally, make sure your training environment includes the horovod package. 如果你使用的是 PyTorch 特选环境,则 horovod 已作为依赖项之一包含在内。If you are using a PyTorch curated environment, horovod is already included as one of the dependencies. 如果使用自己的环境,请确保包含 horovod 依赖项,例如:If you are using your own environment, make sure the horovod dependency is included, for example:

channels:
- conda-forge
dependencies:
- python=3.6.2
- pip:
  - azureml-defaults
  - torch==1.6.0
  - torchvision==0.7.0
  - horovod==0.19.5

若要在 Azure ML 上使用 MPI/Horovod 执行分布式作业,必须指定到 ScriptRunConfig 构造函数的 distributed_job_config 参数的 MpiConfigurationIn order to execute a distributed job using MPI/Horovod on Azure ML, you must specify an MpiConfiguration to the distributed_job_config parameter of the ScriptRunConfig constructor. 以下代码将配置每个节点运行一个进程的 2 节点分布式作业。The below code will configure a 2-node distributed job running one process per node. 如果你还希望每个节点运行多个进程,(即,如果群集 SKU 有多个 GPU),请在 MpiConfiguration 中另外指定 process_count_per_node 参数(默认值为 1)。If you would also like to run multiple processes per node (i.e. if your cluster SKU has multiple GPUs), additionally specify the process_count_per_node parameter in MpiConfiguration (the default is 1).

from azureml.core import ScriptRunConfig
from azureml.core.runconfig import MpiConfiguration

src = ScriptRunConfig(source_directory=project_folder,
                      script='pytorch_horovod_mnist.py',
                      compute_target=compute_target,
                      environment=pytorch_env,
                      distributed_job_config=MpiConfiguration(node_count=2))

有关如何在 Azure ML 上使用 Horovod 运行分布式 PyTorch 的完整教程,请参阅使用 Horovod 的分布式 PyTorchFor a full tutorial on running distributed PyTorch with Horovod on Azure ML, see Distributed PyTorch with Horovod.

DistributedDataParallelDistributedDataParallel

如果你使用的是 PyTorch 的内置 DistributedDataParallel 模块(在训练代码中使用 torch.distributed 包生成),也可以通过 Azure ML 来启动分布式作业。If you are using PyTorch's built-in DistributedDataParallel module that is built using the torch.distributed package in your training code, you can also launch the distributed job via Azure ML.

若要使用 DistributedDataParallel 运行分布式 PyTorch 作业,请将 PyTorchConfiguration 指定到 ScriptRunConfig 构造函数的 distributed_job_config 参数。In order to run a distributed PyTorch job with DistributedDataParallel, specify a PyTorchConfiguration to the distributed_job_config parameter of the ScriptRunConfig constructor. 若要将 NCCL 后端用于 torch.distributed,请在 PyTorchConfiguration 中指定 communication_backend='Nccl'To use the NCCL backend for torch.distributed, specify communication_backend='Nccl' in the PyTorchConfiguration. 以下代码将配置一个 2 节点分布式作业。The below code will configure a 2-node distributed job. 对于 PyTorch 分布式 GPU 训练,建议使用 NCCL 后端作为后端。The NCCL backend is the recommended backend for PyTorch distributed GPU training.

对于通过 PyTorchConfiguration 配置的分布式 PyTorch 作业,Azure ML 将在计算目标的节点上设置以下环境变量:For distributed PyTorch jobs configured via PyTorchConfiguration, Azure ML will set the following environment variables on the nodes of the compute target:

  • AZ_BATCHAI_PYTORCH_INIT_METHOD:进程组的共享文件系统初始化的 URLAZ_BATCHAI_PYTORCH_INIT_METHOD: URL for the shared file system initialization of the process group
  • AZ_BATCHAI_TASK_INDEX:工作进程的全局排名AZ_BATCHAI_TASK_INDEX: global rank of the worker process

可以通过 ScriptRunConfig 的 arguments 参数将这些环境变量指定到训练脚本的相应参数。You can specify these environment variables to the corresponding arguments of the training script via the arguments parameter of ScriptRunConfig.

from azureml.core import ScriptRunConfig
from azureml.core.runconfig import PyTorchConfiguration

args = ['--dist-backend', 'nccl',
        '--dist-url', '$AZ_BATCHAI_PYTORCH_INIT_METHOD',
        '--rank', '$AZ_BATCHAI_TASK_INDEX',
        '--world-size', 2]

src = ScriptRunConfig(source_directory=project_folder,
                      script='pytorch_mnist.py',
                      arguments=args,
                      compute_target=compute_target,
                      environment=pytorch_env,
                      distributed_job_config=PyTorchConfiguration(communication_backend='Nccl', node_count=2))

如果要改用 Gloo 后端来进行分布式训练,请改为指定 communication_backend='Gloo'If you would instead like to use the Gloo backend for distributed training, specify communication_backend='Gloo' instead. 对于分布式 CPU 训练,建议使用 Gloo 后端。The Gloo backend is recommended for distributed CPU training.

有关如何在 Azure ML 上运行分布式 PyTorch 的完整教程,请参阅使用 DistributedDataParallel 的分布式 PyTorchFor a full tutorial on running distributed PyTorch on Azure ML, see Distributed PyTorch with DistributedDataParallel.

导出到 ONNXExport to ONNX

若要使用 ONNX 运行时优化推理,请将训练后的 PyTorch 模型转换为 ONNX 格式。To optimize inference with the ONNX Runtime, convert your trained PyTorch model to the ONNX format. 推理或模型评分是将部署的模型用于预测(通常针对生产数据)的阶段。Inference, or model scoring, is the phase where the deployed model is used for prediction, most commonly on production data. 有关示例,请参阅此教程See the tutorial for an example.

后续步骤Next steps

在本文中,你使用 Azure 机器学习中的 PyTorch 训练并注册了一个深度学习神经网络。In this article, you trained and registered a deep learning, neural network using PyTorch on Azure Machine Learning. 若要了解如何部署模型,请继续参阅模型部署一文。To learn how to deploy a model, continue on to our model deployment article.