优化Pytorch的数据加载
背景
在利用深度学习解决图像问题时,影响训练效率最大的有时候是GPU,有时候也可能是CPU和你的磁盘。 很多设计不当的任务,在训练神经网络的时候,大部分时间都是在从磁盘中读取数据,而不是做 Backpropagation 。这种症状的体现是使用 Nividia-smi 查看 GPU 使用率时,Memory-Usage 占用率很高,但是 GPU-Util 时常为 0% ,如下图所示:
解决方案
如何解决这种问题呢?在 Nvidia 提出的分布式框架 Apex 里面,我们在源码里面找到了一个简单的解决方案: NVIDIA/apex
class data_prefetcher():
def __init__(self, loader):
self.loader = iter(loader)
self.stream = torch.cuda.Stream()
self.mean = torch.tensor([0.485 * 255, 0.456 * 255, 0.406 * 255]).cuda().view(1,3,1,1)
self.std = torch.tensor([0.229 * 255, 0.224 * 255, 0.225 * 255]).cuda().view(1,3,1,1)
# With Amp, it isn't necessary to manually convert data to half.
# if args.fp16:
# self.mean = self.mean.half()
# self.std = self.std.half()
self.preload()
def preload(self):
try:
self.next_input, self.next_target = next(self.loader)
except StopIteration:
self.next_input = None
self.next_target = None
return
with torch.cuda.stream(self.stream):
self.next_input = self.next_input.cuda(non_blocking=True)
self.next_target = self.next_target.cuda(non_blocking=True)
# With Amp, it isn't necessary to manually convert data to half.
# if args.fp16:
# self.next_input = self.next_input.half()
# else:
self.next_input = self.next_input.float()
self.next_input = self.next_input.sub_(self.mean).div_(self.std)
我们能看到 Nvidia 是在读取每次数据返回给网络的时候,预读取下一次迭代需要的数据,那么对我们自己的训练代码只需要做下面的改造:
training_data_loader = DataLoader(
dataset=train_dataset,
num_workers=opts.threads,
batch_size=opts.batchSize,
pin_memory=True,
shuffle=True,
)
for iteration, batch in enumerate(training_data_loader, 1):
# 训练代码
#-------------升级后---------
data, label = prefetcher.next()
iteration = 0
while data is not None:
iteration += 1
# 训练代码
data, label = prefetcher.next()
这样子我们的 Dataloader 就像打了鸡血一样提高了效率很多,如下图:
其他方案
预处理提速
- 尽量减少每次读取数据时的预处理操作,可以考虑把一些固定的操作,例如 resize ,事先处理好保存下来,训练的时候直接拿来用
- Linux上将预处理搬到GPU上加速:NVIDIA/DALI
IO提速
- 使用更快的图片处理库
- opencv一般要比 PIL 要快
- 对于jpeg读取,可以尝试 jpeg4py
- 存bmp图(降低解码时间)
- 小图拼起来存放(降低读取次数,对于大规模的小文件读取,建议转成单独的文件,可以选择的格式:TFRecord(Tensorflow)、recordIO(recordIO)、hdf5、 pth、n5、lmdb 等等 Efficient-PyTorch
- TFRecord
- 借助lmdb数据库: Image2LMDB LMDB PySODToolBox
- 借助内存:直接载到内存里面,或者把把内存映射成磁盘
- 借助固态:把读取速度慢的机械硬盘换成 NVME 固态
训练策略
- 在训练中使用低精度(FP16甚至INT8、二值网络)表示取代原有精度(FP32)表示
- NVIDIA/Apex使用
- torch.backends.cudnn.benchmark = True
- Do numpy-like operations on the GPU wherever you can
- Free up memory using del
- Avoid unnecessary transfer of data from the GPU
- Use pinned memory, and use non_blocking=False to parallelize data transfer and GPU number crunching
模型设计
来自ShuffleNetV2的结论:
内存访问消耗时间,memory access cost 缩写为 MAC。
- 卷积层输入输出通道一致:卷积层的输入和输出特征通道数相等时MAC最小,此时模型速度最快
- 减少卷积分组:过多的group操作会增大MAC,从而使模型速度变慢
- 减少模型分支:模型中的分支数量越少,模型速度越快
- 减少element-wise操作:element-wise操作所带来的时间消耗远比在FLOPs上的体现的数值要多,因此要尽可能减少element-wise操作(depthwise convolution也具有低FLOPs、高MAC的特点)
其他:
- 降低复杂度:例如模型裁剪和剪枝,减少模型层数和参数规模
- 改模型结构:例如模型蒸馏,通过知识蒸馏方法来获取小模型
本文来源: