文档HOT
下载
中文
注册

数据加载优化

在PyTorch模型中，数据加载部分的逻辑一般是DataLoader及其衍生类，在DataLoader加载数据中，要注意以下两个核心点：
  • 第一个是数据加载的预处理部分，数据的预处理通常会写在datasets里，数据的预处理包括对文本、图片、视频和语音等不同格式数据的处理，数据预处理耗时长是比较容易识别出来的，一般通过打点计时。
  • 第二个就是要确定好数据读取的方式，一般而言，每张卡都去读取数据是比较理想的，如果存在0卡读取，广播给其他卡的数据读取形式，要关注其性能如何，很多时候都会让模型性能严重劣化。

硬件类

  • 检查数据存放的硬盘，最好在NVME（Non-Volatile Memory Express，一种更快、更可靠的存储设备，适合存储需要高速读写的数据）盘上，当硬盘为共享存储时，需要注意IO瓶颈。
  • 检查NUMA（Non-Uniform Memory Access，非一致性内存访问）数量，可以通过lscpu命令查看，一般需要NUMA数量为4或者8。

参数类

DataLoader入参：
  • num_workers是一个比较常用的参数，一般而言，对于图像或视频等复杂数据，可以将num_workers设置的大一些（如4或8），通过多进程并行提高数据读入和解析的速度。在使用这个参数的时候，需要结合实际情况，如视频解析本身并不存在瓶颈，过多的num_workers则会增加进程开销，因此需要酌情考虑。
  • 在内存允许的情况下，保持pin_memory=True，在PyTorch中，该参数可以写为pin_memory_device。
  • 一般而言，可以加上persistent_workers=True，该设置可以减小进程销毁或申请的开销，不过这个方法也可能带来内存瓶颈。

代码技巧类

  • 在数据预处理时间较长而影响模型训练时，可以尝试预取数据，即Data prefetcher。
  • Infinite DataLoader同样可以缓解epoch间加载缓慢的问题。
  • 提前缓存数据或制作二进制数据，轻量化数据加载，若原始数据是压缩或编码格式，则会引发解压缩或解码，加重CPU负载，容易引起host bound问题，建议提前处理成二进制格式，快速进行数据读取。
  • 对于NPU预处理数据比较慢，例如dlrm模型，可以考虑将预处理放到datasets里用CPU多num_workers处理数据。
  • 对于句子之间长度差距比较大的句子，组成batch时可以考虑resample一下。对于NLP任务，存在输入的sentence长短不一，tokenize之后需要padding成相同长度才能组成一个batch的情况。对于一些数据集中，所有sentence长度差别很大的情况，如果随机取句子，那么需要补的padding就很多，造成很多冗余的计算量。此时，可以通过调整DataSampler，尽量将长度接近的sentence组成一个batch，这样造成很多冗余的计算量就可以减少很多，训练整体的吞吐量会有显著提升。缺点：可能稍微会影响数据的随机性。
  • collate_fn用于将多个workers从datasets中读取的batch_list组合成一个batch，这个组合过程是可以加速的，例如原来每个workers都要对各自的数据transpose并且减均值除方差再合并成一个batch。此时可以先合并成一个batch，再放到Device上，并且用Device算力完成transpose和除方差操作，达到减轻CPU负载的目的。以此原理，timm仓库设计了fast_collate函数，已经是timm训练的标配，能够有效提升性能。