深度学习模型的训练/推理过程涉及很多步骤。在有限的时间和资源条件下，每个迭代的速度越快，整个模型的预测性能就越快。我收集了几个PyTorch技巧，以最大化内存使用效率和最小化运行时间。为了更好地利用这些技巧，我们还需要理解它们如何以及为什么有效。

我首先提供一个完整的列表和一些代码片段，这样你就可以开始优化你的脚本了。然后我一个一个地详细地研究它们。对于每个技巧，我还提供了代码片段和注释，告诉你它是特定于设备类型(CPU/GPU)还是模型类型。

列表:

  • 数据加载

    1、把数据放到SSD中

    2、Dataloader(dataset, num_workers=4*num_GPU)

    3、Dataloader(dataset, pin_memory=True)

  • 数据操作

    4、直接在设备中创建torch.Tensor，不要在一个设备中创建再移动到另一个设备中

    5、避免CPU和GPU之间不必要的数据传输

    6、使用torch.from_numpy(numpy_array)或者torch.as_tensor(others)

    7、在数据传输操作可以重叠时，使用tensor.to(non_blocking=True)

    8、使用PyTorch JIT将元素操作融合到单个kernel中。

  • 模型结构

    9、在使用混合精度的FP16时，对于所有不同架构设计，设置尺寸为8的倍数

  • 训练

    10、将batch size设置为8的倍数，最大化GPU内存的使用

    11、前向的时候使用混合精度（后向的使用不用）

    12、在优化器更新权重之前，设置梯度为None，model.zero_grad(set_to_none=True)

    13、梯度积累：每隔x个batch更新一次权重，模拟大batch size的效果

  • 推理/验证

    14、关闭梯度计算

  • CNN (卷积神经网络) 特有的

    15、torch.backends.cudnn.benchmark = True

    16、对于4D NCHW Tensors，使用channels_last的内存格式

    17、在batch normalization之前的卷积层可以去掉bias

  • 分布式

    18、用DistributedDataParallel代替DataParallel

内容中包含的图片若涉及版权问题，请及时与我们联系删除