训练模型时 遇到速度过慢时的深思 & 速度提升 （From GPU & CPU）

GPU

查看GPU使用情况 配置

nvidia_smi

单机多卡并行训练

torch.nn.DataParallel

在使用多GPU的时候, 使用os.environ[‘CUDA_VISIBLE_DEVICES’]来限制使用的GPU个数
例如我要使用第0和第3编号的GPU, 那么只需要在程序中设置:

os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'

Note:
这个参数的设定要保证在模型加载到gpu上之前, 一般在程序开始的时候就设定好这个参数

之后, 将模型加载到多GPU上面

如果是模型:

model = nn.DataParallel(model)
model = model.cuda()

如果是数据:

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

pytorch官网给的示例代码

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)

DataParallel的内部代码:

class DataParallel(Module):
    def __init__(self, module, device_ids=None, output_device=None, dim=0):
        super(DataParallel, self).__init__()

        if not torch.cuda.is_available():
            self.module = module
            self.device_ids = []
            return

        if device_ids is None:
            device_ids = list(range(torch.cuda.device_count()))
        if output_device is None:
            output_device = device_ids[0]

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

深入考究
使用os.environ['CUDA_VISIBLE_DEVICES']对可以使用的显卡进行限定之后, 显卡的实际编号和程序看到的编号应该是不一样的

例如
上面我们设定的是os.environ['CUDA_VISIBLE_DEVICES']="0,2",但是程序看到的显卡编号应该被改成了'0,1', 也就是说程序所使用的显卡编号实际上是经过了一次映射之后才会映射到真正的显卡编号上面的,
例如
这里的程序看到的1对应实际的2

平衡DataParallel带来的显存使用不平衡的问题

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

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

2. 然后我将其中的平衡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()

torch.nn.parallel.DistributedDataParallel

pytorch的官网建议使用DistributedDataParallel来代替DataParallel

因为DistributedDataParallel比DataParallel运行的更快

然后 显存分屏的更加均衡
且DistributedDataParallel 功能更加强悍
例如
分布式的模型(一个模型太大, 以至于无法放到一个GPU上运行, 需要分开到多个GPU上面执行).
只有DistributedDataParallel支持分布式的模型像单机模型那样可以进行多机多卡的运算.

先设定好os.environ[‘CUDA_VISIBLE_DEVICES’]
然后再进行下面的步骤.

因为DistributedDataParallel是支持多机多卡的, 所以这个需要先初始化一下:

torch.distributed.init_process_group(backend='nccl', init_method='tcp://localhost:23456', rank=0, world_size=1)

第一个参数是pytorch支持的通讯后端, 这里单机多卡, 这个就是走走过场.
第二个参数是各个机器之间通讯的方式，这里是单机多卡, 设置成localhost, 后面的端口找一个空着没用的就OK.
**rank**是标识主机和从机, 这里就一个主机, 设置成0就OK.
**world_size**是标识使用几个主机, 这里就一个主机, 设置成1就OK, 设置多了代码不允许.

其实如果是使用单机多卡的情况下, 根据pytorch的官方代码distributeddataparallel, 是直接可以使用下面的代码的:

torch.distributed.init_process_group(backend="nccl")
model = DistributedDataParallel(model) # device_ids will include all GPU devices by default

Note:
如果使用这句代码, 直接在pycharm或者别的编辑器中,是没法正常运行的
因为这个需要在shell的命令行中运行, 如果想要正确执行这段代码, 假设这段代码的名字是main.py, 可以使用如下的方法进行(参考1 参考2):

python -m torch.distributed.launch main.py

如果使用了argparse, 一定要在参数里面加上–local_rank, 否则运行还是会出错的

之后就和使用DataParallel很类似了.

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

Note: 这里要先将model加载到GPU 然后才能使用DistributedDataParallel进行分发
之后的使用和DataParallel就基本一样了

多机多gpu训练

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

Reference

pytorch/examples/imagenet/main.py

Distributed-VGG-F

使用半精度训练

更好的显卡，更轻的模型

batch_size

增大 batch size 提高epoch速度
但是收敛速度也会变慢
需要再适当升高学习率

CPU

查看占用率
top命令

top -bn 1 -i -c

data loader

线程数 arg 增加
多线程加载数据
dataloader 的param.

num_works=4

如果数据集很小，使用多线程加载数据可能会更慢，因为多线程有一定的开销。
在这种情况下，最好使用单线程读取数据。

减少日志IO操作频率

文件的IO操作，这会导致GPU得不到连续性使用，整体速度特别慢。

使用pin_memory和num_workers

总结

有的时候模型训练慢并不是因为显卡不行或者模型太大，而是在跑模型过程中有一些其他的操作导致速度很慢
如：IO操作

1. GPU 升级
2. 硬盘读取，重点观察数据是否量大且容量较小
3. data loader 线程数arg 增加
4. 观察CPU占用率
5. 内存占用率
6. IO速度等

Reference

https://zhuanlan.zhihu.com/p/86441879