高性能PyTorch训练

参考：
https://mp.weixin.qq.com/s/foB44Fm-IhX3yaawn_aZgg

数据预处理

几乎每个训练管道都以 Dataset 类开始。它负责提供数据样本。任何必要的数据转换和扩充都可能在此进行。简而言之，Dataset 能报告其规模大小以及在给定索引时，给出数据样本。

如果你要处理类图像的数据（2D、3D 扫描），那么磁盘 I/O 可能会成为瓶颈。为了获取原始像素数据，你的代码需要从磁盘中读取数据并解码图像到内存。每个任务都是迅速的，但是当你需要尽快处理成百上千或者成千上万个任务时，可能就成了一个挑战。像 NVidia 这样的库会提供一个 GPU 加速的 JPEG 解码。如果你在数据处理管道中遇到了 IO 瓶颈，这种方法绝对值得一试。

还有另外一个选择，SSD 磁盘的访问时间约为 0.08–0.16 毫秒。RAM 的访问时间是纳秒级别的。我们可以直接将数据存入内存。

建议 1：如果可能的话，将数据的全部或部分移至 RAM。
如果你的内存中有足够多的 RAM 来加载和保存你的训练数据，这是从管道中排除最慢的数据检索步骤最简单的方法。

class RAMDataset(Dataset):
  def __init__(image_fnames, targets):
    self.targets = targets
    self.images = []
    for fname in tqdm(image_fnames, desc="Loading files in RAM"):
      with open(fname, "rb") as f:
        self.images.append(f.read())

  def __len__(self):
    return len(self.targets)

  def __getitem__(self, index):
    target = self.targets[index]
    image, retval = cv2.imdecode(self.images[index], cv2.IMREAD_COLOR)
    return image, target

可以将每个文件的二进制内容保持不变，并在 RAM 中进行即时解码，或者对未压缩的图像进行讲解码，并保留原始像素。

---

---

参考：
https://zhuanlan.zhihu.com/p/450 强烈推荐

Data模块

1. python图像处理用的最多的两个库是opencv和Pillow（PIL），但是两者读取出来的图像并不一样，opencv读取的图像格式的三个通道是BGR形式的，但是PIL是RGB格式的。这个问题看起来很小，但是衍生出来的坑可以有很多，最常见的场景就是数据增强和预训练模型中。比如有些数据增强的方法是基于channel维度的，比如megengine里面的HueTransform，在这一行代码显然是需要确保图像是BGR的，但是经常会有人只看有Transform就无脑用了，从来没有考虑过这些问题。

2. 接上条，RGB和BGR的另一个问题就是导致预训练模型载入后训练的方式不对，最常见的场景就是预训练模型的input channel是RGB的（例如torch官方来的预训练模型），然后你用cv2做数据处理，最后还忘了convert成RGB的格式，那么就是会有问题。这个问题应该很多炼丹的同学没有注意过，我之前写CenterNet-better就发现CenterNet存在这么一个问题，要知道当时这可是一个有着3k多star的仓库，但是从来没有人意识到有这个问题。当然，依照我的经验，如果你训练的iter足够多，即使你的channel有问题，对于结果的影响也会非常小。不过，既然能做对，为啥不注意这些问题一次性做对呢？

3. torchvision中提供的模型，都是输入图像经过了ToTensor操作train出来的。也就是说最后在进入网络之前会统一除以255从而将网络的输入变到0到1之间。torchvision的文档给出了他们使用的mean和std，也是0-1的mean和std。如果你使用torch预训练的模型，但是输入还是0-255的，那么恭喜你，在载入模型上你又会踩一个大坑（要么你的图像先除以255，要么你的code中mean和std的数值都要乘以255）。

4. ToTensor之后接数据处理的坑。上一条说了ToTensor之后图像变成了0到1的，但是一些数据增强对数值做处理的时候，是针对标准图像，很多人ToTensor之后接了这样一个数据增强，最后就是练出来的丹是废的（心疼电费QaQ）。

5. 数据集里面有一个图特别诡异，只要train到那一张图就会炸显存（CUDA OOM），别的图训练起来都没有问题，应该怎么处理？通常出现这个问题，首先判断数据本身是不是有问题。如果数据本身有问题，在一开始生成Dataset对象的时候去掉就行了。如果数据本身没有问题，只不过因为一些特殊原因导致显存炸了（比如检测中图像的GT boxes过多的问题），可以catch一个CUDA OOM的error之后将一些逻辑放在CPU上，最后retry一下，这样只是会慢一个iter，但是训练过程还是可以完整走完的，在我们开源的YOLOX里有类似的参考code。

6. pytorch中dataloader的坑。有时候会遇到pytorch num_workers=0（也就是单进程）没有问题，但是多进程就会报一些看不懂的错的现象，这种情况通常是因为torch到了ulimit的上限，更核心的原因是torch的dataloader不会释放文件描述符（参考issue）。可以ulimit -n 看一下机器的设置。跑程序之前修改一下对应的数值。

7. opencv和dataloader的神奇联动。很多人经常来问为啥要写cv2.setNumThreads(0)，其实是因为cv2在做resize等op的时候会用多线程，当torch的dataloader是多进程的时候，多进程套多线程，很容易就卡死了（具体哪里死锁了我没探究很深）。除了setNumThreads之外，通常还要加一句cv2.ocl.setUseOpenCL(False)，原因是cv2使用opencl和cuda一起用的时候通常会拖慢速度，加了万事大吉，说不定还能加速。

8. dataloader会在epoch结束之后进行类似重新加载的操作，复现这个问题的code稍微有些长，放在后面了。这个问题算是可以说是一个高级bug/feature了，可能导致的问题之一就是炼丹师在本地的code上进行了一些修改，然后训练过程直接加载进去了。解决方法也很简单，让你的sampler源源不断地产生数据就好，这样即使本地code有修改也不会加载进去。

Module模块

1. BatchNorm在训练和推断的时候的行为是不一致的。这也是新人最常见的错误（类似的算子还有dropout，这里提一嘴，pytorch的dropout在eval的时候行为是Identity，之前有遇到过实习生说dropout加了没效果，直到我看了他的code： x = F.dropout(x, p=0.5)

2. BatchNorm叠加分布式训练的坑。在使用DDP（DistributedDataParallel）进行训练的时候，每张卡上的BN统计量是可能不一样的，仔细检查broadcast_buffer这个参数。DDP的默认行为是在forward之前将rank0 的 buffer做一次broadcast（broadcast_buffer=True），但是一些常用的开源检测仓库是将broadcast_buffer设置成False的（参考：mmdet 和 detectron2，我猜是在检测任务中因为batchsize过小，统一用卡0的统计量会掉点）这个问题在一边训练一边测试的code中更常见，比如说你train了5个epoch，然后要分布式测试一下。一般的逻辑是将数据集分到每块卡上，每块卡进行inference，最后gather到卡0上进行测点。但是因为每张卡统计量是不一样的，所以和那种把卡0的模型broadcast到不同卡上测试出来的结果是不一样的。这也是为啥通常训练完测的点和单独起了一个测试脚本跑出来的点不一样的原因（当然你用SyncBN就不会有这个问题）。

3. Pytorch的SyncBN在1.5之前一直实现的有bug，所以有一些老仓库是存在使用SyncBN之后掉点的问题的。

4. 用了多卡开多尺度训练，明明尺度更小了，但是速度好像不是很理想？这个问题涉及到多卡的原理，因为分布式训练的时候，在得到新的参数之后往往需要进行一次同步。假设有两张卡，卡0的尺度非常小，卡1的尺度非常大，那么就会出现卡0始终在等卡1，于是就出现了虽然有的尺度变小了，但是整体的训练速度并没有变快的现象（木桶效应）。解决这个问题的思路就是尽量把负载拉均衡一些。

5. 多卡的小batch模拟大batch（梯度累积）的坑。假设我们在单卡下只能塞下batchsize = 2，那么为了模拟一个batchsize = 8的效果，通常的做法是forward / backward 4次，不清理梯度，step一次（当然考虑BN的统计量问题这种做法和单纯的batchsize=8肯定还是有一些差别的）。在多卡下，因为调用loss.backward的时候会做grad的同步，所以说前三次调用backward的时候需要加ddp.no_sync的context manager（不加的话，第一次bp之后，各个卡上的grad此时会进行同步），最后一次则不需要加。当然，我看很多仓库并没有这么做，我只能理解他们就是单纯想做梯度累积（BTW，加了ddp.no_sync会使得程序快一些，毕竟加了之后bp过程是无通讯的）。

6. 浮点数的加法其实不遵守交换律的，这个通常能衍生出来GPU上的运算结果不能严格复现的现象。

训练模块

1. FP16训练/混合精度训练。使用Apex训练混合精度模型，在保存checkpoint用于继续训练的时候，除了model和optimizer本身的state_dict之外，还需要保存一下amp的state_dict，这个在amp的文档中也有提过。（当然，经验上来说忘了保存影响不大，会多花几个iter search一个loss scalar出来）

2. 多机分布式训练卡死的问题。好友 @NoahSYZhang 遇到的一个坑。场景是先申请了两个8卡机，然后机器1和机器2用前4块卡做通讯（local rank最大都是4，总共是两机8卡）。可以初始化process group，但是在使用DDP的时候会卡死。原因在于pytorch在做DDP的时候会猜测一个rank，参考code。对于上面的场景，第二个机器上因为存在卡5到卡8，而对应的rank也是5到8，所以DDP就会认为自己需要同步的是卡5到卡8，于是就卡死了。

3. 在使用AMP的时候，使用Adam/AdamW优化器之后NaN，前面的iter没有任何异常现象。通常是optimizer里面的eps的问题，调整一下eps的数值就好了（比如1e-3），因为默认的eps是1e-8，在fp16下浮点运算容易出NaN

4. 梯度为0 和 参数是否更新 没有必然关系。因为grad并不是最终的参数更新量，最终的参数更新量是在optimizer里面进行计算的。一个最简单的例子就是设置了weight decay不为0，当optimizer的weight decay不为0 的时候，最终的参数更新量都会加上 lr * wd * param ，所以 grad为0并不等价于参数量不会更新。一些可以refer的code（此处以megengine的code为例，pytorch只是把逻辑写成了cpp来加速）