超全Pytorch多GPU训练

本文为多篇博客的综合，原博客的链接在博客最后的位置。

---

一、DataParallel使用方式

1. 使用方法

使用多卡训练的方式有很多，当然前提是我们的设备中存在两个及以上的GPU：使用命令nvidia-smi查看当前Ubuntu平台的GPU数量(Windows平台类似)，其中每个GPU被编上了序号：[0,1]：

在我们设备中确实存在多卡的条件下，可以采用pytorch官方的代码：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
    model = nn.DataParallel(model)

model.to(device)

如果我们机子中有很多显卡(例如我们有八张显卡)，但我们只想使用0、2号显卡，那么我们可以：

net = torch.nn.DataParallel(model, device_ids=[0, 2])
net = net.cuda()

或者这样：

os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, [0,2]))
net = torch.nn.DataParallel(model)
net = net.cuda()
# CUDA_VISIBLE_DEVICES 表示当前可以被python环境程序检测到的显卡

如果我们不设定好要使用的device_ids的话, 程序会自动找到这个机器上面可以用的所有的显卡, 然后用于训练. 但是因为我们前面使用os.environ['CUDA_VISIBLE_DEVICES']限定了这个程序可以使用的显卡, 所以这个地方程序如果自己获取的话, 获取到的其实就是我们上面设定的那几个显卡。

没有进行深入得到考究, 但是感觉使用os.environ['CUDA_VISIBLE_DEVICES']对可以使用的显卡进行限定之后, 显卡的实际编号和程序看到的编号应该是不一样的, 例如上面我们设定的是os.environ['CUDA_VISIBLE_DEVICES']="0,2", 但是程序看到的显卡编号应该被改成了’0,1’, 也就是说程序所使用的显卡编号实际上是经过了一次映射之后才会映射到真正的显卡编号上面的, 例如这里的程序看到的1对应实际的2。

数据迁移用以下代码即可：

inputs = inputs.cuda()
labels = labels.cuda()

2. 负载不均衡（显存使用不平衡）

这个问题其实讨论的也比较多了, 官方给的解决方案就是使用 DistributedDataParallel来代替 DataParallel(实际上DistributedDataParallel显存分配的也不是很平衡), 但是从某些角度来说, DataParallel使用起来确实比较方便, 而且最近使用 DistributedDataParallel 遇到一些小问题. 所以这里提供一个解决显存使用不平衡问题的方案:

首先这次的解决方案来自transformer-XL的官方代码: https://github.com/kimiyoung/transformer-xl

然后将其中的平衡GPU显存的代码提取了出来(原代码好像有点小问题)放到了github上面:https://github.com/Link-Li/Balanced-DataParallel

这里的代码是原作者继承了 DataParallel 类之后进行了改写:

class BalancedDataParallel(DataParallel):
    def __init__(self, gpu0_bsz, *args, **kwargs):
        self.gpu0_bsz = gpu0_bsz
        super().__init__(*args, **kwargs)
    ...

这个 BalancedDataParallel 类使用起来和 DataParallel 类似, 下面是一个示例代码:

my_net = MyNet()
my_net = BalancedDataParallel(gpu0_bsz // acc_grad, my_net, dim=0).cuda()

这里包含三个参数, 第一个参数是第一个GPU要分配多大的batch_size, 但是要注意, 如果你使用了梯度累积, 那么这里传入的是每次进行运算的实际batch_size大小. 举个例子, 比如你在3个GPU上面跑代码, 但是一个GPU最大只能跑3条数据, 但是因为0号GPU还要做一些数据的整合操作, 于是0号GPU只能跑2条数据, 这样一算, 你可以跑的大小是2+3+3=8, 于是你可以设置下面的这样的参数:

batch_szie = 8
gpu0_bsz = 2
acc_grad = 1
my_net = MyNet()
my_net = BalancedDataParallel(gpu0_bsz // acc_grad, my_net, dim=0).cuda()

这个时候突然想跑个batch size是16的怎么办呢, 那就是4+6+6=16了, 这样设置累积梯度为2就行了:

batch_szie = 16
gpu0_bsz = 4
acc_grad = 2
my_net = MyNet()
my_net = BalancedDataParallel(gpu0_bsz // acc_grad, my_net, dim=0).cuda()

二、DistributedParallel使用方式

在前言中提到过，另一种方法DistributedParallel，虽然主要的目标为分布式训练，但也是可以实现单主机多GPU方式训练的，只不过比上一种方法稍微麻烦一点，但是训练速度和效果比上一种更好。

为什么呢？

请看官方相关介绍：

nccl backend is currently the fastest and highly recommended backend to be used with Multi-Process Single-GPU distributed training and this applies to both single-node and multi-node distributed training

1. 单机多卡

好了，来说说具体的使用方法(下面展示一个node也就是一个主机的情况)为：

python -m torch.distributed.launch --nproc_per_node=你的GPU数量
               YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other
               arguments of your training script)

上述的命令和我们平常的命令稍有区别，这里我们用到了torch.distributed.launch这个module，我们选择运行的方式变换为python -m，上面相当于使用torch.distributed.launch.py去运行我们的YOUR_TRAINING_SCRIPT.py，其中torch.distributed.launch会向我们的运行程序传递一些变量。

为此，我们的YOUR_TRAINING_SCRIPT.py也就是我们的训练代码中这样写(省略多余代码，只保留核心代码)：

import torch.distributed as dist
# 这个参数是torch.distributed.launch传递过来的，我们设置位置参数来接受，local_rank代表当前程序进程使用的GPU标号
parser.add_argument("--local_rank", type=int, default=0) 

def synchronize():
    """
    Helper function to synchronize (barrier) among all processes when
    using distributed training
    """
    if not dist.is_available():
        return
    if not dist.is_initialized():
        return
    world_size = dist.get_world_size()
    if world_size == 1:
        return
    dist.barrier()


## WORLD_SIZE 由torch.distributed.launch.py产生 具体数值为 nproc_per_node*node(主机数，这里为1)
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1

is_distributed = num_gpus > 1

if is_distributed:
    torch.cuda.set_device(args.local_rank)  # 这里设定每一个进程使用的GPU是一定的
    torch.distributed.init_process_group(
        backend="nccl", init_method="env://"
    )
    synchronize()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 以Darknet为例
model = Darknet(opt.model_def).to(device)
# 将模型移至到DistributedDataParallel中，此时就可以进行训练了
if is_distributed:
    model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[args.local_rank], output_device=args.local_rank,
        # this should be removed if we update BatchNorm stats
        broadcast_buffers=True,
    )
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
	model = model.module

# 数据操作
dataset = Dataset()
if is_distributed:
	train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
else:
	train_sampler = None

dataloader = torch.utils.data.DataLoader(
	dataset,
	# 省略中间参数
	sampler=train_sampler,
)

我们要注意，上述代码在运行的过程中产生了很多个，具体多少个取决你GPU的数量，这也是为什么上面需要torch.cuda.set_device(args.local_rank)设定默认的GPU，因为torch.distributed.launch为我们触发了n个YOUR_TRAINING_SCRIPT.py进程，n就是我们将要使用的GPU数量。

命令行执行方式：

python -m torch.distributed.launch --nproc_per_node=2 train.py [--args 你的参数列表]

有一点想问的，我们每次必须要使用命令行的方式去运行吗？当然也可以一键解决，如果我们使用Pycharm，只需要配置成下面这样就可以了：

2. 多机多卡

在单机多gpu可以满足的情况下, 绝对不建议使用多机多gpu进行训练, 我经过测试, 发现多台机器之间传输数据的时间非常慢, 主要是因为我测试的机器可能只是千兆网卡, 再加上别的一些损耗, 网络的传输速度跟不上, 导致训练速度实际很慢. 我看一个github上面的人说在单机8显卡可以满足的情况下, 最好不要进行多机多卡训练

建议看这两份代码, 实际运行一下, 才会真的理解怎么使用

pytorch/examples/imagenet/main.py

Distributed-VGG-F

2.1.初始化

初始化操作一般在程序刚开始的时候进行

在进行多机多gpu进行训练的时候, 需要先使用torch.distributed.init_process_group()进行初始化. torch.distributed.init_process_group()包含四个常用的参数

backend: 后端, 实际上是多个机器之间交换数据的协议
init_method: 机器之间交换数据, 需要指定一个主节点, 而这个参数就是指定主节点的
world_size: 介绍都是说是进程, 实际就是机器的个数, 例如两台机器一起训练的话, world_size就设置为2
rank: 区分主节点和从节点的, 主节点为0, 剩余的为了1-(N-1), N为要使用的机器的数量, 也就是world_size

2.1.1.初始化backend

首先要初始化的是backend, 也就是俗称的后端, 在pytorch的官方教程中提供了以下这些后端

根据官网的介绍, 如果是使用cpu的分布式计算, 建议使用gloo, 因为表中可以看到 gloo对cpu的支持是最好的, 然后如果使用gpu进行分布式计算, 建议使用nccl, 实际测试中我也感觉到, 当使用gpu的时候, nccl的效率是高于gloo的. 根据博客和官网的态度, 好像都不怎么推荐在多gpu的时候使用mpi

对于后端选择好了之后, 我们需要设置一下网络接口, 因为多个主机之间肯定是使用网络进行交换, 那肯定就涉及到ip之类的, 对于nccl和gloo一般会自己寻找网络接口, 但是某些时候, 比如我测试用的服务器, 不知道是系统有点古老, 还是网卡比较多, 需要自己手动设置. 设置的方法也比较简单, 在Python的代码中, 使用下面的代码进行设置就行:

import os
# 以下二选一, 第一个是使用gloo后端需要设置的, 第二个是使用nccl需要设置的
os.environ['GLOO_SOCKET_IFNAME'] = 'eth0'
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'

我们怎么知道自己的网络接口呢, 打开命令行, 然后输入ifconfig, 然后找到那个带自己ip地址的就是了, 我见过的一般就是em0, eth0, esp2s0之类的, 当然具体的根据你自己的填写. 如果没装ifconfig, 输入命令会报错, 但是根据报错提示安装一个就行了.

2.1.2.初始化init_method

初始化init_method的方法有两种, 一种是使用TCP进行初始化, 另外一种是使用共享文件系统进行初始化

2.1.2.1.使用TCP初始化

看代码:

import torch.distributed as dist

dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456',
                        rank=rank, world_size=world_size)

注意这里使用格式为tcp://ip:端口号, 首先ip地址是你的主节点的ip地址, 也就是rank参数为0的那个主机的ip地址, 然后再选择一个空闲的端口号, 这样就可以初始化init_method了.

2.1.2.2.使用共享文件系统初始化

好像看到有些人并不推荐这种方法, 因为这个方法好像比TCP初始化要没法, 搞不好和你硬盘的格式还有关系, 特别是window的硬盘格式和Ubuntu的还不一样, 我没有测试这个方法, 看代码:

import torch.distributed as dist

dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile',
                        rank=rank, world_size=world_size)

根据官网介绍, 要注意提供的共享文件一开始应该是不存在的, 但是这个方法又不会在自己执行结束删除文件, 所以下次再进行初始化的时候, 需要手动删除上次的文件, 所以比较麻烦, 而且官网给了一堆警告, 再次说明了这个方法不如TCP初始化的简单.

2.1.3.初始化rank和world_size

这里其实没有多难, 你需要确保, 不同机器的rank值不同, 但是主机的rank必须为0, 而且使用init_method的ip一定是rank为0的主机, 其次world_size是你的主机数量, 你不能随便设置这个数值, 你的参与训练的主机数量达不到world_size的设置值时, 代码是不会执行的.

2.1.4.初始化中一些需要注意的地方

首先是代码的统一性, 所有的节点上面的代码, 建议完全一样, 不然有可能会出现一些问题, 其次, 这些初始化的参数强烈建议通过argparse模块(命令行参数的形式)输入, 不建议写死在代码中, 也不建议使用pycharm之类的IDE进行代码的运行, 强烈建议使用命令行直接运行.

其次是运行代码的命令方面的问题, 例如使用下面的命令运行代码distributed.py:

python distributed.py -bk nccl -im tcp://10.10.10.1:12345 -rn 0 -ws 2

上面的代码是在主节点上运行, 所以设置rank为0, 同时设置了使用两个主机, 在从节点运行的时候, 输入的代码是下面这样:

python distributed.py -bk nccl -im tcp://10.10.10.1:12345 -rn 1 -ws 2

一定要注意的是, 只能修改rank的值, 其他的值一律不得修改, 否则程序就卡死了初始化到这里也就结束了.

2.2.数据的处理-DataLoader

其实数据的处理和正常的代码的数据处理非常类似, 但是因为多机多卡涉及到了效率问题, 所以这里才会使用torch.utils.data.distributed.DistributedSampler来规避数据传输的问题. 首先看下面的代码:

print("Initialize Dataloaders...")
# Define the transform for the data. Notice, we must resize to 224x224 with this dataset and model.
transform = transforms.Compose(
    [transforms.Resize(224),
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# Initialize Datasets. STL10 will automatically download if not present
trainset = datasets.STL10(root='./data', split='train', download=True, transform=transform)
valset = datasets.STL10(root='./data', split='test', download=True, transform=transform)

# Create DistributedSampler to handle distributing the dataset across nodes when training
# This can only be called after torch.distributed.init_process_group is called
# 这一句就是和平时使用有点不一样的地方
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)

# Create the Dataloaders to feed data to the training and validation steps
train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=(train_sampler is None), num_workers=workers, pin_memory=False, sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(valset, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=False)

其实单独看这段代码, 和平时写的很类似, 唯一不一样的其实就是这里先将trainset送到了DistributedSampler中创造了一个train_sampler, 然后在构造train_loader的时候, 参数中传入了一个sampler=train_sampler. 使用这些的意图是, 让不同节点的机器加载自己本地的数据进行训练, 也就是说进行多机多卡训练的时候, 不再是从主节点分发数据到各个从节点, 而是各个从节点自己从自己的硬盘上读取数据.

当然了, 如果直接让各个节点自己读取自己的数据, 特别是在训练的时候经常是要打乱数据集进行训练的, 这样就会导致不同的节点加载的数据混乱, 所以这个时候使用DistributedSampler来创造一个sampler提供给DataLoader, sampler的作用自定义一个数据的编号, 然后让DataLoader按照这个编号来提取数据放入到模型中训练, 其中sampler参数和shuffle参数不能同时指定, 如果这个时候还想要可以随机的输入数据, 我们可以在DistributedSampler中指定shuffle参数, 具体的可以参考官网的api, 拉到最后就是DistributedSampler

2.3.模型的处理

模型的处理其实和上面的单机多卡没有多大区别, 还是下面的代码, 但是注意要提前想把模型加载到gpu, 然后才可以加载到DistributedDataParallel

model = model.cuda()
model = nn.parallel.DistributedDataParallel(model)

2.4.模型的保存与加载

这里引用pytorch官方教程的一段代码:

def demo_checkpoint(rank, world_size):
    setup(rank, world_size)

    # setup devices for this process, rank 1 uses GPUs [0, 1, 2, 3] and
    # rank 2 uses GPUs [4, 5, 6, 7].
    n = torch.cuda.device_count() // world_size
    device_ids = list(range(rank * n, (rank + 1) * n))

    model = ToyModel().to(device_ids[0])
    # output_device defaults to device_ids[0]
    ddp_model = DDP(model, device_ids=device_ids)

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    CHECKPOINT_PATH = tempfile.gettempdir() + "/model.checkpoint"
    if rank == 0:
        # All processes should see same parameters as they all start from same
        # random parameters and gradients are synchronized in backward passes.
        # Therefore, saving it in one process is sufficient.
        torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)

    # Use a barrier() to make sure that process 1 loads the model after process
    # 0 saves it.
    dist.barrier()
    # configure map_location properly
    rank0_devices = [x - rank * len(device_ids) for x in device_ids]
    device_pairs = zip(rank0_devices, device_ids)
    map_location = {'cuda:%d' % x: 'cuda:%d' % y for x, y in device_pairs}
    ddp_model.load_state_dict(
        torch.load(CHECKPOINT_PATH, map_location=map_location))

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(device_ids[0])
    loss_fn = nn.MSELoss()
    loss_fn(outputs, labels).backward()
    optimizer.step()

    # Use a barrier() to make sure that all processes have finished reading the
    # checkpoint
    dist.barrier()

    if rank == 0:
        os.remove(CHECKPOINT_PATH)

    cleanup()

我并没有实际操作, 因为多卡多GPU代码运行起来实在是难受, 一次实验可能就得好几分钟, 要是搞错一点可能就得好几十分钟都跑不起来, 最主要的是还要等能用的GPU. 不过看上面的代码, 最重要的实际是这句 dist.barrier(), 这个是来自torch.distributed.barrier(), 根据pytorch的官网的介绍, 这个函数的功能是同步所有的进程, 直到整组(也就是所有节点的所有GPU)到达这个函数的时候, 才会执行后面的代码, 看上面的代码, 可以看到, 在保存模型的时候, 是只找rank为0的点保存模型, 然后在加载模型的时候, 首先得让所有的节点同步一下, 然后给所有的节点加载上模型, 然后在进行下一步的时候, 还要同步一下, 保证所有的节点都读完了模型. 虽然我不清楚这样做的意义是什么, 但是官网说不这样做会导致一些问题, 我并没有实际操作, 不发表意见.

至于保存模型的时候, 是保存哪些节点上面的模型, pytorch推荐的是rank=0的节点, 然后我看在论坛上, 有人也会保存所有节点的模型, 然后进行计算, 至于保存哪些, 我并没有做实验, 所以并不清楚到底哪种最好.

三、注意点

多GPU固然可以提升我们训练的速度，但弊端还有有一些的，有几个我们需要注意的点：

多个GPU的数量尽量为偶数，奇数的GPU有可能会出现中断的情况
选取与GPU数量相适配的数据集，多显卡对于比较小的数据集来说反而不如单个显卡训练的效果好
多GPU训练的时候注意机器的内存是否足够(一般为使用显卡显存x2)，如果不够，建议关闭pin_memory(锁页内存)选项。
采用DistributedDataParallel多GPUs训练的方式比DataParallel更快一些，如果你的Pytorch编译时有nccl的支持，那么最好使用DistributedDataParallel方式。
关于什么是锁页内存：

pin_memory就是锁页内存，创建DataLoader时，设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。
主机中的内存，有两种存在方式，一是锁页，二是不锁页，锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘），而不锁页内存在主机内存不足时，数据会存放在虚拟内存中。显卡中的显存全部是锁页内存,当计算机的内存充足的时候，可以设置pin_memory=True。当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False。

相关资料：
https://blog.csdn.net/tfcy694/article/details/83270701
https://discuss.pytorch.org/t/what-is-the-disadvantage-of-using-pin-memory/1702

总结一句就是，如果机子的内存比较大，建议开启pin_memory=Ture，如果开启后发现有卡顿现象或者内存占用过高，此时建议关闭。

参考文章

• https://pytorch.org/docs/stable/notes/cuda.html#cuda-nn-dataparallel-instead
• https://github.com/facebookresearch/maskrcnn-benchmark
• https://oldpan.me/archives/pytorch-to-use-multiple-gpus
• https://zhuanlan.zhihu.com/p/86441879