赞
踩
在优化PyTorch代码性能的过程中,发现了官网的这篇 Performance Tuning Guide1,觉得很有帮助。因此,在这一边翻译一边加入个人的实践经验供参考。
性能调优指南是一套优化方法和最佳实践方式,可以加速 PyTorch 深度学习模型的训练和推理 。所介绍的技术通常只需改变几行代码就可以实现,并且可以应用于所有领域的各种深度学习模型。
torch.utils.data.DataLoader
支持独立工作子进程中的异步数据加载和数据扩充。DataLoader
的默认设置为num_workers=0
。这意味着数据加载是同步的,在主进程中完成。因此,主训练过程必须等待数据可用后才能继续执行。
设置num_workers>0
,可以支持异步数据加载以及训练与数据加载之间的重叠进行。num_workers
应该根据工作负载、CPU、GPU 和训练数据的位置进行调优。
DataLoader
接受pin_memory
参数,其默认值为False
。当使用 GPU 时,最好设置pin_memory=True
,指示DataLoader
使用固定内存,使内存更快地从主机异步复制到GPU上。
PyTorch 会从需要梯度的张量的所有运算中保存中间缓冲区。通常,验证或推理不需要梯度。torch.no_grad()
上下文管理器可以用来禁用指定代码块内的梯度计算。这将加快运行的速度并减少所需的内存。torch.no_grad()
也可以用作函数的装饰器。
torch.nn.Conv2d()
有bias
参数,默认值为True
(Conv1d
和Conv3d
也是如此)。
如果一个nn.Conv2d
层后面直接连了一个nn.BatchNorm2d
层,那么就不需要卷积中的偏置,改用nn.Conv2d(..., bias=False, ....)
。不需要偏置是因为在第一步中,BatchNorm
减去了平均值,有效地抵消了偏置的影响。
只要BatchNorm
(或其他归一化层)在与卷积偏置的相同维度上进行归一化,这个方法同样适用于一维和三维卷积。
torchvision 提供的模型已经实现了这种优化。
parameter.grad = None
替代model.zero_grad()
或optimizer.zero_grad()
不再调用
model.zero_grad()
# or
optimizer.zero_grad()
使梯度归零,改用以下方法:
for param in model.parameters():
param.grad = None
第二个代码段不会将每个参数的内存清零,而且后续的反向传播将使用赋值而不是加法来存储梯度,这减少了内存运算的数量。
将梯度设置为None
与设置为零的数值表现略有不同,详情请参阅文档2。
或者,从 PyTorch 1.7 开始,可以调用model
或 optimizer.zero_grad(set_to_none=True)
。
实测 PyTorch 2.0.0 set_to_none
默认为True
。
逐点运算(逐元素加法、乘法、数学函数 - sin()
、cos()
、sigmoid()
等)可以被融合到一个内核中,以分摊内存访问时间和内核启动时间。
PyTorch JIT 可以自动融合内核。尽管可能有更多的融合机会尚未在编译器中实现,而且并非所有的设备类型都得到同等支持。(如果不了解内核融合,可以参见3)
逐点运算是受内存约束的,对于每个运算,PyTorch都会启动一个单独的内核。每个内核从内存中加载数据,执行计算(这一步通常不会消耗太多资源),并将结果存储回内存中。
融合的算子对于多个融合逐点运算只启动一个内核,并且只向内存加载/存储一次数据。这使得 JIT 对于激活函数、优化器、自定义 RNN 单元等非常有用。
在最简单的情况下,可以通过对函数定义应用torch.jit.script
装饰器来实现融合,例如:
@torch.jit.script
def fused_gelu(x):
return x * 0.5 * (1.0 + torch.erf(x / 1.41421))
关于更多高级用例,请参考 TorchScript 文档4。
PyTorch 1.5 为卷积网络引入了对channels_last
内存格式的支持。这种格式是为了与 AMP(Automatic Mixed Precision 自动混合精度)一起使用,以进一步加速带有 Tensor Cores 的卷积神经网络。
对channels_last
的支持是实验性的,但它有望用于标准的计算机视觉模型(例如ResNet-50、SSD)。要将模型转换为channels_last
格式,请遵循 Channels Last 内存格式教程5。该教程包括关于转换现有模型的部分。
缓冲区检查点是一种减轻模型训练的内存容量负担的技术。它并非存储所有层的输入来计算反向传播中的上游梯度,而是存储少数层的输入,其他层则在反向传播中重新计算。内存需求的减少可以增加批处理大小,从而提高利用率。
应谨慎选择检查点的目标。最好不要存储再计算开销较小的大层的输出。目标层的例子是激活函数(如ReLU
、Sigmoid
、Tanh
),上/下采样和积累深度较小的矩阵-向量运算。
PyTorch 支持原生的torch.utils.checkpoint
API来自动执行检查点和再计算。
如果使用 Huggingface Transformers 的模型,可以调用函数model.gradient_checkpointing_enable()
来实现。
许多 PyTorch API 是用于调试的,应在常规训练运行中禁用:
torch.autograd.detect_anomaly
或torch.autograd.set_detect_anomaly(True)
torch.autograd.profiler.emit_nvtx
,torch.autograd.profiler.profile
torch.autograd.gradcheck
或torch.autograd.gradgradcheck
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。