- 混合精度训练:FP32量化为FP16
- 大Batch训练:对前几次梯度进行累加,然后统一进行参数更新,从而变相实现大Batch训练
- 梯度检查点:训练时间换显存,在前向阶段传递到 checkpoint 中的 forward 函数会以 torch.no_grad 模式运行,并且仅仅保存输入参数和 forward 函数,在反向阶段重新计算其 forward 输出值。
- 尽量实现向量化运算(即使开的数组多占用空间或者操作步数多),
- 使用
@torch.jit.script
装饰器会将pointwise操作融合成一个kernel来执行,比如激活函数(pointwise操作)中的多个运算 - pin_memory=True,减少数据传送时间
- 直接在GPU上创建tensor,eg:
tensor = torch.rand([10,5], device=torch.device('cuda:0'))
torch.tensor()
始终复制数据,numpy到tensor使用torch.from_numpy(numpy_array)
更快,tensor(或list、tuple)到tensor使用torch.as_tensor(others)
更快- 使用
tensor.to(device, non_blocking=True)
允许异步来overlap data transfers and kernel execution,数据流向构成一个DAG - 输入、输出、中间参数的维度尽量是8的倍数
- 如果输入大小不经常变化,在训练前设置
torch.backends.cudnn.benchmark=True
,cuDNN会自动寻找最好的算法来计算卷积(因为cuDNN算法计算size不同的卷积的性能不同) - 在BN层之前的卷积层中不使用bias(因为BN第一步会减去均值)
- Use channels_last memory format for 4D NCHW Tensors
- use
DistributedDataParallel
instead ofDataParallel
参考资料:
推荐:Optimize PyTorch Performance for Speed and Memory Effciency(2022)