数据加载优化

在PyTorch模型中,数据加载部分的逻辑一般是DataLoader及其衍生类,在DataLoader加载数据中,要注意以下两个核心点:

硬件类

参数类

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

代码技巧类