点击下方卡片，关注“FightingCV”公众号
回复“AI”即可获得超100G人工智能的教程
点击进入→FightingCV交流群
来源：江大白
导读
PyTorch是当前人工智能比较成熟和广泛应用的框架。那么PyTorch中一些常见的小技巧发夹是否熟悉呢？本文分享了PyTorch中的18个小技巧，希望能对各位读者有所帮助。
深度学习模型的训练/推理过程涉及很多步骤。在有限的时间和资源条件下，每个迭代的速度越快，整个模型的预测性能就越快。我收集了几个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=True16、对于4D NCHW Tensors，使用channels_last的内存格式
17、在batch normalization之前的卷积层可以去掉bias
分布式
18、用
DistributedDataParallel代替DataParallel
#CombiningthetipsNo.7,11,12,13:nonblocking,AMP,setting
#gradientsasNone,andlargereffectivebatchsize
model.train()
#ResetthegradientstoNone
optimizer.zero_grad(set_to_none=True)
scaler=GradScaler()
fori,(features,target)inenumerate(dataloader):
#thesetwocallsarenonblockingandoverlapping
features=features.to('cuda:0',non_blocking=True)
target=target.to('cuda:0',non_blocking=True)
#Forwardpasswithmixedprecision
withtorch.cuda.amp.autocast():#autocastasacontextmanager
output=model(features)
loss=criterion(output,target)
#Backwardpasswithoutmixedprecision
#It'snotrecommendedtousemixedprecisionforbackwardpass
#Becauseweneedmorepreciseloss
scaler.scale(loss).backward()
#Onlyupdateweightseveryother2iterations
#Effectivebatchsizeisdoubled
if(i+1)%2==0or(i+1)==len(dataloader):
#scaler.step()firstunscalesthegradients.
#IfthesegradientscontaininfsorNaNs,
#optimizer.step()isskipped.
scaler.step(optimizer)
#Ifoptimizer.step()wasskipped,
#scalingfactorisreducedbythebackoff_factor
#inGradScaler()
scaler.update()
#ResetthegradientstoNone
optimizer.zero_grad(set_to_none=True)
总的来说，你可以通过3个关键点来优化时间和内存使用。首先，尽可能减少i/o(输入/输出)，使模型管道更多的用于计算，而不是用于i/o(带宽限制或内存限制)。这样，我们就可以利用GPU及其他专用硬件来加速这些计算。第二，尽量重叠过程，以节省时间。第三，最大限度地提高内存使用效率，节约内存。然后，节省内存可以启用更大的batch size大小，从而节省更多的时间。拥有更多的时间有助于更快的模型开发周期，并导致更好的模型性能。
有些机器有不同的硬盘驱动器，如HHD和SSD。建议将项目中使用的数据移动到SSD(或具有更好i/o的硬盘驱动器)以获得更快的速度。
num_workers=0使数据加载需要在训练完成后或前一个处理已完成后进行。设置num_workers>0有望加快速度，特别是对于大数据的i/o和增强。具体到GPU，有实验发现num_workers=4*num_GPU具有最好的性能。也就是说，你也可以为你的机器测试最佳的num_workers。需要注意的是，高num_workers将会有很大的内存消耗开销，这也是意料之中的，因为更多的数据副本正在内存中同时处理。
Dataloader(dataset,num_workers=4*num_GPU)
GPU无法直接从CPU的可分页内存中访问数据。设置pin_memory=True可以为CPU主机上的数据直接分配临时内存，节省将数据从可分页内存转移到临时内存(即固定内存又称页面锁定内存)的时间。该设置可以与num_workers=4*num_GPU结合使用。
Dataloader(dataset,pin_memory=True)
只要你需要torch.Tensor，首先尝试在要使用它们的设备上创建它们。不要使用原生Python或NumPy创建数据，然后将其转换为torch.Tensor。在大多数情况下，如果你要在GPU中使用它们，直接在GPU中创建它们。
#Randomnumbersbetween0and1
#Sameasnp.random.rand([10,5])
tensor=torch.rand([10,5],device=torch.device('cuda:0'))
#Randomnumbersfromnormaldistributionwithmean0andvariance1
#Sameasnp.random.randn([10,5])
tensor=torch.randn([10,5],device=torch.device('cuda:0'))
唯一的语法差异是NumPy中的随机数生成需要额外的random，例如：np.random.rand()vstorch.rand()。许多其他函数在NumPy中也有相应的函数：
torch.empty(),torch.zeros(),torch.full(),torch.ones(),torch.eye(),torch.randint(),torch.rand(),torch.randn()
正如我在指导思想中提到的，我们希望尽可能地减少I/O。注意下面这些命令：
#BAD!AVOIDTHEMIFUNNECESSARY!
print(cuda_tensor)
cuda_tensor.cpu()
cuda_tensor.to_device('cpu')
cpu_tensor.cuda()
cpu_tensor.to_device('cuda')
cuda_tensor.item()
cuda_tensor.numpy()
cuda_tensor.nonzero()
cuda_tensor.tolist()
#PythoncontrolflowwhichdependsonoperationresultsofCUDAtensors
if(cuda_tensor!=0).all():
run_func()
torch.tensor()会拷贝数据
如果源设备和目标设备都是CPU，torch.from_numpy和torch.as_tensor不会创建数据拷贝。如果源数据是NumPy数组，使用torch.from_numpy(numpy_array)会更快。如果源数据是一个具有相同数据类型和设备类型的张量，那么torch.as_tensor(others)可以避免拷贝数据。others可以是Python的list，tuple，或者torch.tensor。如果源设备和目标设备不同，那么我们可以使用下一个技巧。
torch.from_numpy(numpy_array)
torch.as_tensor(others)
本质上，non_blocking=True允许异步数据传输以减少执行时间。
forfeatures,targetinloader:
#thesetwocallsarenonblockingandoverlapping
features=features.to('cuda:0',non_blocking=True)
target=target.to('cuda:0',non_blocking=True)
#Thisisasynchronizationpoint
#Itwillwaitforprevioustwolines
output=model(features)
点操作包括常见的数学操作，通常是内存受限的。PyTorch JIT会自动将相邻的点操作融合到一个内核中，以保存多次内存读/写操作。例如，通过将5个核融合成1个核，gelu函数可以被加速4倍。
@torch.jit.script#JITdecorator
deffused_gelu(x):
returnx*0.5*(1.0+torch.erf(x/1.41421))
为了最大限度地提高GPU的计算效率，最好保证不同的架构设计(包括神经网络的输入输出尺寸/维数/通道数和batch size大小)是8的倍数甚至更大的2的幂(如64、128和最大256)。这是因为当矩阵的维数与2的幂倍数对齐时，Nvidia gpu的张量核心(Tensor Cores)在矩阵乘法方面可以获得最佳性能。矩阵乘法是最常用的操作，也可能是瓶颈，所以它是我们能确保张量/矩阵/向量的维数能被2的幂整除的最好方法(例如，8、64、128，最多256)。
这些实验显示设置输出维度和batch size大小为8的倍数，比如(33712、4088、4096)相比33708，batch size为4084或者4095这些不能被8整除的数可以加速计算1.3倍到 4倍。加速度大小取决于过程类型(例如，向前传递或梯度计算)和cuBLAS版本。特别是，如果你使用NLP，请记住检查输出维度，这通常是词汇表大小。
使用大于256的倍数不会增加更多的好处，但也没有害处。这些设置取决于cuBLAS和cuDNN版本以及GPU架构。你可以在文档中找到矩阵维数的特定张量核心要求。由于目前PyTorch AMP多使用FP16，而FP16需要8的倍数，所以通常推荐使用8的倍数。如果你有更高级的GPU，比如A100，那么你可以选择64的倍数。如果你使用的是AMD GPU，你可能需要检查AMD的文档。
除了将batch size大小设置为8的倍数外，我们还将batch size大小最大化，直到它达到GPU的内存限制。这样，我们可以用更少的时间来完成一个epoch。
有些操作不需要float64或float32的精度。因此，将操作设置为较低的精度可以节省内存和执行时间。对于各种应用，英伟达报告称具有Tensor Cores的GPU的混合精度可以提高3.5到25倍的速度。
值得注意的是，通常矩阵越大，混合精度加速度越高。在较大的神经网络中(例如BERT)，实验表明混合精度可以加快2.75倍的训练，并减少37%的内存使用。具有Volta, Turing, Ampere或Hopper架构的较新的GPU设备(例如，T4, V100, RTX 2060, 2070, 2080, 2080 Ti, A100, RTX 3090, RTX 3080，和RTX 3070)可以从混合精度中受益更多，因为他们有Tensor Core架构，它相比CUDA cores有特殊的优化。
值得一提的是，采用Hopper架构的H100预计将于2022年第三季度发布，支持FP8 (float8)。PyTorch AMP可能会支持FP8(目前v1.11.0还不支持FP8)。
在实践中，你需要在模型精度性能和速度性能之间找到一个最佳点。我之前确实发现混合精度可能会降低模型的精度，这取决于算法，数据和问题。
使用自动混合精度(AMP)很容易在PyTorch中利用混合精度。PyTorch中的默认浮点类型是float32。AMP将通过使用float16来进行一组操作(例如，matmul,linear,conv2d)来节省内存和时间。AMP会自动cast到float32的一些操作(例如，mse_loss,softmax等)。有些操作(例如add)可以操作最宽的输入类型。例如，如果一个变量是float32，另一个变量是float16，那么加法结果将是float32。
autocast自动应用精度到不同的操作。因为损失和梯度是按照float16精度计算的，当它们太小时，梯度可能会“下溢”并变成零。GradScaler通过将损失乘以一个比例因子来防止下溢，根据比例损失计算梯度，然后在优化器更新权重之前取消梯度的比例。如果缩放因子太大或太小，并导致inf或NaN，则缩放因子将在下一个迭代中更新缩放因子。
scaler=GradScaler()
forfeatures,targetindata:
#Forwardpasswithmixedprecision
withtorch.cuda.amp.autocast():#autocastasacontextmanager
output=model(features)
loss=criterion(output,target)
#Backwardpasswithoutmixedprecision
#It'snotrecommendedtousemixedprecisionforbackwardpass
#Becauseweneedmorepreciseloss
scaler.scale(loss).backward()
#scaler.step()firstunscalesthegradients.
#IfthesegradientscontaininfsorNaNs,
#optimizer.step()isskipped.
scaler.step(optimizer)
#Ifoptimizer.step()wasskipped,
#scalingfactorisreducedbythebackoff_factorinGradScaler()
scaler.update()
你也可以使用autocast作为前向传递函数的装饰器。
classAutocastModel(nn.Module):
...
@autocast()#autocastasadecorator
defforward(self,input):
x=self.model(input)
returnx
通过model.zero_grad()或optimizer.zero_grad()将对所有参数执行memset，并通过读写操作更新梯度。但是，将梯度设置为None将不会执行memset，并且将使用“只写”操作更新梯度。因此，设置梯度为None更快。
#Resetgradientsbeforeeachstepofoptimizer
forparaminmodel.parameters():
param.grad=None
#or(PyTorch>=1.7)
model.zero_grad(set_to_none=True)
#or(PyTorch>=1.7)
optimizer.zero_grad(set_to_none=True)
这个技巧是关于从更多的数据样本积累梯度，以便对梯度的估计更准确，权重更新更接近局部/全局最小值。这在batch size较小的情况下更有帮助(由于GPU内存限制较小或每个样本的数据量较大)。
fori,(features,target)inenumerate(dataloader):
#Forwardpass
output=model(features)
loss=criterion(output,target)
#Backwardpass
loss.backward()
#Onlyupdateweightseveryother2iterations
#Effectivebatchsizeisdoubled
if(i+1)%2==0or(i+1)==len(dataloader):
#Updateweights
optimizer.step()
#ResetthegradientstoNone
optimizer.zero_grad(set_to_none=True)
实际上，如果只计算模型的输出，那么梯度计算对于推断和验证步骤并不是必需的。PyTorch使用一个中间内存缓冲区来处理requires_grad=True变量中涉及的操作。因此，如果我们知道不需要任何涉及梯度的操作，通过禁用梯度计算来进行推断/验证，就可以避免使用额外的资源。
#torch.no_grad()asacontextmanager:
withtorch.no_grad():
output=model(input)
#torch.no_grad()asafunctiondecorator:
@torch.no_grad()
defvalidation(model,input):
output=model(input)
returnoutput
在训练循环之前设置torch.backends.cudnn.benchmark=True可以加速计算。由于计算不同内核大小卷积的cuDNN算法的性能不同，自动调优器可以运行一个基准来找到最佳算法。当你的输入大小不经常改变时，建议开启这个设置。如果输入大小经常改变，那么自动调优器就需要太频繁地进行基准测试，这可能会损害性能。它可以将向前和向后传播速度提高1.27x到1.70x。
torch.backends.cudnn.benchmark=True
使用channels_last内存格式以逐像素的方式保存图像，作为内存中最密集的格式。原始4D NCHW张量在内存中按每个通道(红/绿/蓝)顺序存储。转换之后，x=x.to(memory_format=torch.channels_last)，数据在内存中被重组为NHWC (channels_last格式)。你可以看到RGB层的每个像素更近了。据报道，这种NHWC格式与FP16的AMP一起使用可以获得8%到35%的加速。
目前，它仍处于beta测试阶段，仅支持4D NCHW张量和一组模型(例如，alexnet，mnasnet家族，mobilenet_v2，resnet家族，shufflenet_v2，squeezenet1，vgg家族)。但我可以肯定，这将成为一个标准的优化。
N,C,H,W=10,3,32,32
x=torch.rand(N,C,H,W)
#Strideisthegapbetweenoneelementtothenextone
#inadimension.
print(x.stride())
#(3072,1024,32,1)#ConvertthetensortoNHWCinmemory
x2=x.to(memory_format=torch.channels_last)
print(x2.shape)#(10,3,32,32)asdimensionsorderpreserved
print(x2.stride())#(3072,1,96,3),whicharesmaller
print((x==x2).all())#Truebecausethevalueswerenotchanged
这是可行的，因为在数学上，bias可以通过batch normalization的均值减法来抵消。我们可以节省模型参数、运行时的内存。
nn.Conv2d(...,bias=False)
对于多GPU来说，即使只有单个节点，也总是优先使用DistributedDataParallel而不是DataParallel，因为DistributedDataParallel应用于多进程，并为每个GPU创建一个进程，从而绕过Python全局解释器锁(GIL)并提高速度。
在这篇文章中，我列出了一个清单，并提供了18个PyTorch技巧的代码片段。然后，我逐一解释了它们在不同方面的工作原理和原因，包括数据加载、数据操作、模型架构、训练、推断、cnn特定的优化和分布式计算。一旦你深入理解了它们的工作原理，你可能会找到适用于任何深度学习框架中的深度学习建模的通用原则。
往期回顾
基础知识
【CV知识点汇总与解析】|损失函数篇
【CV知识点汇总与解析】|激活函数篇
【CV知识点汇总与解析】| optimizer和学习率篇
【CV知识点汇总与解析】| 正则化篇
【CV知识点汇总与解析】| 参数初始化篇
【CV知识点汇总与解析】| 卷积和池化篇 （超多图警告）
最新论文解析
CVPR2022 | Attention机制是为了找最相关的item？中科大团队反其道而行之！
ECCV2022 Oral | SeqTR：一个简单而通用的 Visual Grounding网络
如何训练用于图像检索的Vision Transformer？Facebook研究员解决了这个问题！
ICLR22 Workshop | 用两个模型解决一个任务，意大利学者提出维基百科上的高效检索模型
See Finer, See More！腾讯&上交提出IVT，越看越精细，进行精细全面的跨模态对比！
MM2022｜兼具低级和高级表征，百度提出利用显式高级语义增强视频文本检索
MM2022 | 用StyleGAN进行数据增强，真的太好用了
MM2022 | 在特征空间中的多模态数据增强方法
ECCV2022|港中文MM Lab证明Frozen的CLIP 模型是高效视频学习者
ECCV2022｜只能11%的参数就能优于Swin，微软提出快速预训练蒸馏方法TinyViT
CVPR2022|比VinVL快一万倍！人大提出交互协同的双流视觉语言预训练模型COTS，又快又好！
CVPR2022 Oral｜通过多尺度token聚合分流自注意力，代码已开源
CVPR Oral | 谷歌&斯坦福（李飞飞组）提出TIRG，用组合的文本和图像来进行图像检索
