how to fully utilize gpu using pytorch?

deep learning, machine learning, pytorch

写在前面

我们不妨先考虑 单 gpu 的情况，torch 的 dataloader 非常不稳定，经常出现无法拉满 gpu / gpu util 不断震荡的情况，这是因为 gpu 总是在等待 dataloader 输送数据，造成了很大的 latency。解决这个问题我们首先要搞清楚 bound 在哪里，并且根据合适的方法来因地制宜。

在一个正常的 pytorch 训练过程中，我们通常会重载 Dataset 来实现自己的数据读取，并通过 Dataloader 接口来进行数据的 streaming read. 注意 dataloader 均利用 cpu 进行数据读取与 data augmentation, 我们还需要调用 _.cuda() 来将数据拷贝至 gpu，然后 gpu 才能用数据进行 forward 与 backwork process. 上述过程中，第一个可能的 bound 来自 IO, dataset 通常需要从磁盘上读取数据；第二个 bound 来自 cpu, cpu 读取图片，解码图片的速度是否足够快；第三个 bound 来自 cpu 与 gpu 的数据拷贝过程；

IO bound?

第一是 IO bound, 即把数据从磁盘读入内存的时间开销。

显然，更好的硬盘可以带来更快的 IO, 条件允许的情况下可以将数据都放在 ssd 硬盘上；再进一步，如果数据都已经在内存中，IO 开销自然消失。如果我们的数据量不大，比如 nyuv2 就可以直接预先读入内存，而不需要每一个 epoch 都重新从磁盘中读取；再者，linux 系统可以将数据拷贝至 memory mapped disk(没有仔细考究用词)，比如 /dev/shm 下，或者自己挂载一块内存区域 mount -t tmpfs -o size=20g tmpfs /userhome/memory_data

以上的方法将 IO 读取看作一个包裹好的接口, 其实我们还可以更深一步地思考。即便是相同的数据，IO 读取速度也会有差异，这和数据实际存储的方式有关，对于这样一块需要反复访问的数据，如果在底层能够分配一块区域来存储，必然能提高读取速度，这也是 lmdb, h5py 等库能加速读取的原理。

cpu bound?

第二是 cpu bound，即读取数据的过程以及解码数据的过程

以图片来举例，我通常利用 opencv 来读取 cv2.imread()（其他库大同小异）, opencv 首先将数据读入内存，然后再解码并编码成 numpy.ndarray。第一个过程，IO 层面的加速已在上一节中描述，在 cpu 这里，我们还可以利用多线程/进程的技术进一步加速读取，由于 python 自身的限制，torch 可以利用 multi-process 来加速读取，我们需要指定 dataloader workers 的数量，来确定读取的进程数。需要注意的是线程不是越多越好，有条件的话可以通过实验决定一个相对较快又不影响正常系统运行的线程数。第二个过程，想要加速我们可以利用其他的数据格式，比如上文提到的 h5py, lmdb 以及 TFRecord 都是将数据重新编码，从而加快读取的速度。

cpu -> gpu bound?

第三是 cpu 到 gpu copy bound

通常来说这个过程不会成为 bound, 但是仍有几处加速空间。首先是 dataloader 的 pin_memory 参数，该参数为 True 时，将会把 tensor 拷贝至一个 page-locked memory buffer, 这样的话 gpu 取数据总是在一个 memory 的固定 buffer 读取，从而加速 transfer 的速度。一点需要注意的是，使用 cumstom collect_fn 的时候需要自己完成 pin_memmory 的逻辑，详见 torch docs on memory pinning 当 pin_memory=True 后，我们还可以利用异步 gpu 读取来进一步加快训练速度，只需要在 _.cuda() 或 _.to() 时设置 non_blocking=True 即可。此时 gpu 的数据 copy 与 computation 异步进行，computation 也就不用等待 copy，二者可以 overlap, 减少各自的 latency. 理想情况下，gpu 无需等待数据的输入过程，buffer 中总有提前拷贝的数据，可以不断地进行 computation, 此时 gpu 的 utilization 自然就上去了；

what if we get rid of cpu?

当整个训练过程 bound 在 gpu computation 之外时，gpu utilization 就不高，整个系统的潜力就被浪费了，如果我们尝试的各种方法都不能解决问题的时候，还有一个路子：nvidia 的 dali dataloader. DALI 可以将上述提到的大部分过程都转移到 gpu 上高效操作（读取，decode, augmentation 等），从而跳脱出 cpu 的限制，在通常情况下都能大大提高训练速度.

multi-gpu

我们在上面讨论了单 gpu 的情况，当涉及到多 gpu 时，我们还需考虑各个 gpu 的负载均匀。通常来讲，在 parameter server 模式下会有一个主 gpu 来收集各个 gpu 的梯度信息来更新模型，这也是 tensorflow 与 tensorpack 经常用的模式。torch 的多 gpu 训练非常简单，仅仅需要 DataParallel wrap 一层即可，具体参见 torch official doc；在使用 DataParallel 之后，主 gpu 的 util 会高一些，但是总体来说各个 gpu 基本负载均衡，如果出现严重的不平衡，肯定是模型或者模型的某部分没有正确的部署在 gpu 上。DataDistributed 本来是用于分布式训练的接口，近期很多 post 都提到了利用这个接口，即便在单机训练上也能有更好的表现。

what if gpu bound

在已经喂饱 gpu 的情况下，我们还可以通过混合精度的方法来进一步通过模型 precision 的损失来加速训练，apex 可以做到这一点，我暂时没有尝试过，有兴趣的朋友可以试试。

miscs

sources

TODOs

  • h5py, lmdb, 图片不同的解码格式