利用PyTorch-Ignite轻松实现分布式训练

562 阅读6分钟

PyTorch-Ignite logo

作者François Cokelaer,Priyansi,Sylvain Desroziers,Victor Fomin

编写支持不同平台、硬件配置(GPU、TPU)和通信框架的不可知 分布式代码是很繁琐的。在这篇博客中,我们将讨论PyTorch-Ignite如何以最小的代码改动解决这个问题。

内容

前提条件

本博客假定您对以下方面有一定的了解。

  1. PyTorch的分布式包、它所提供的后端集体功能。在这篇博客中,我们将重点讨论分布式数据并行代码

  2. PyTorch-Ignite。请参考本博客,以获得快速的高层次概述。

简介

PyTorch-Ignite的 ignite.distributed(idist)子模块在v0.4.0版本(2020年7月)中引入,可以快速将单进程代码变成其数据分布式版本。

因此,你现在可以在所有支持的后端无缝地运行同一版本的代码。

在这篇博文中,我们将比较PyTorch-Ignite的API和Torch native的分布式代码,并强调前者的差异和易用性。我们还将展示Ignite的auto_* 方法如何自动使你的代码与上述分布式后端兼容,这样你就只需要带上自己的模型、优化器和数据加载器对象。

代码片段以及运行所有脚本的命令,都在一个单独的资源库中提供。

然后,我们还将介绍通过torch本地torch.multiprocessing.spawn ,以及通过多个分布式启动器催生进程的几种方式,以强调Pytorch-Ignite的idist ,特别是在不改变代码的情况下,如何处理它。

更多关于发射器实验的信息可以在这里找到。

🔥 Pytorch-Ignite统一分布式API

我们需要为不同的分布式后端编写不同的代码。这可能很繁琐,特别是如果你想在不同的硬件配置上运行你的代码。Pytorch-Ignite的idist ,由于高层次的帮助方法,将为你完成所有的工作。

🔍 关注帮助性的auto_* 方法。

这个方法为非分布式和可用的分布式配置调整了逻辑。下面是分布式模型实例化的等效代码片断。

PyTorch-IgnitePyTorch DDP
HorovodTorch XLA

此外,它还兼容NVIDIA/apex

model

Torch原生AMP

model

这个方法可以无缝地适应非分布式和可用的分布式配置的优化器逻辑。以下是分布式优化器实例化的等效代码片断。

PyTorch-IgnitePyTorch DDP
HorovodTorch XLA

该方法为目标设备上的非分布式和可用的分布式配置无缝适应数据加载逻辑。

此外,auto_dataloader() ,根据分布式配置上下文自动调整批次大小,形成了在多个设备上加载样本批次的通用方法。

以下是分布式数据加载步骤的等效代码片断。

PyTorch-IgnitePyTorch DDP
HorovodTorch XLA

此外,idist 提供了集体操作,如all_reduce,all_gather, 和broadcast ,可以与所有支持的分布式框架一起使用。请参阅我们的文档以了解更多细节。

例子

下面的代码片段强调了每个分布式后端在相同使用情况下与idist API相比的API的特殊性。PyTorch本地代码可用于DDP、Horovod和XLA/TPU设备。

PyTorch-Ignite的统一代码片段可以与标准的PyTorch后端如gloonccl ,也可以与Horovod和XLA的TPU设备一起运行。请注意,该代码不那么冗长,但是,用户仍然可以完全控制训练循环。

下面的例子是介绍性的。对于一个使用PyTorch-Ignite的更强大的生产级例子,请参考这里

这些实验的完整源代码可以在这里找到。

PyTorch-Ignite - Torch原生分布式数据并行 - Horovod - XLA/TPU

|

PyTorch-Ignite

|

PyTorch DDP

| | --- | --- | | [

源代码

](github.com/pytorch-ign…) | [

源代码

](github.com/pytorch-ign…) | | | |

浩沃德

|

火炬XLA

| | [

源代码

](github.com/pytorch-ign…) | [

源代码

](github.com/pytorch-ign…) | | | |

注意

你也可以将idist 与其他分布式API混合使用,如下所示。

dist

运行分布式代码

PyTorch-Ignite的idist ,也统一了分布式代码的启动方法,并通过ignite.distributed.launcher.Parallel(idist Parallel)上下文管理器使分布式配置的设置更加简单。

这个上下文管理器有能力生成nproc_per_node (作为脚本参数传递)的子进程,并根据提供的后台初始化处理组,或者使用像torch.distributed.launch,slurm,horovodrun 这样的工具,只以一般的方式初始化给出backend 参数的处理组。

torch.multiprocessing.spawn

在这种情况下,idist Parallel 正在使用引擎盖下的本地火炬torch.multiprocessing.spawn 方法,以便运行分布式配置。这里nproc_per_node ,作为一个spawn参数被传递。

# Running with gloo

使用分布式启动器

PyTorch-Ignite的idist Parallel 上下文管理器也与多个分布式启动器兼容。

使用Torch.distributed.launch

这里我们使用torch.distributed.launch 脚本,以便催生进程。

python -m torch.distributed.launch --nproc_per_node 

使用horovodrun

horovodrun -np 

注意

为了运行这个例子,并避免安装过程,你可以从PyTorch-Ignite的docker镜像中抽取一个预装了Horovod的镜像。它将包括Horovod与gloo 控制器和nccl 支持。

docker run --gpus all -it -v 

使用slurm

使用slurm ,不需要对代码进行任何修改,就可以达到同样的效果。

srun --nodes

或者使用sbatch script.bash 与脚本文件script.bash

#!/bin/bash

结束语

通过上面的例子我们看到,管理分布式计算的多种配置和规格从未如此简单。只需几行字,我们就可以在任何地方并行化和执行代码,同时保持控制和简单性。

参考文献

  • idist-snippets:本文章中使用的完整代码。

  • why-ignite:分布式数据并行的例子:本地pytorch,pytorch-ignite,slurm。

  • CIFAR10的例子,在CIFAR10上用多种配置进行分布式训练。1个或多个GPU,多个节点和GPU,TPU。

接下来的步骤

  • 如果你想了解更多关于PyTorch-Ignite的信息或有任何进一步的疑问,这里有我们的GitHub文档Discord

  • PyTorch-Ignite目前是由一个志愿者团队维护的,我们正在寻找更多的贡献者。请参阅CONTRIBUTING.md以了解您的贡献方式。

  • 请在TwitterFacebook上关注我们,了解所有PyTorch-Ignite的最新动态。