PyTorch模型优化思路¶
前言¶
近年来AI模型成为ASC竞赛选题的香饽饽，如2022年的三四两题分别是基于pytorch的大型语言模型，和基于tensorflow的分子动力学模型。AI模型的优化可以从许多方面入手，本文参考Lorenze Kuhn的文章Here are 17 ways of making PyTorch training faster – what did I miss?，对其进行翻译和总结，介绍pytorch框架的通用优化思路。
注意：本文所讲方法仅供参考，帮助大家打开思路。你可能会发现这些方法并不会有太大效果，甚至没有效果。如果想有所突破，必须靠自己深度阅读代码，对症下药。
考虑换一种学习率¶
学习率对模型收敛速度和泛化能力有很大影响。Cyclical Learning Rates 和oneCycle learning rate 是 Leslie N. Smith提出的两种自适应学习率方法。（参见这里和这里）oneCycle learning rate schedule看起来像这样：
PyTorch实现了这两种方法:torch.optim.lr_scheduler.CyclicLR 和torch.optim.lr_scheduler.OneCycleLR。具体用法请见官方文档。
注意这两种方法需要给定超参数。关于超参数的选择请看这条帖子和这个仓库
在DataLoader里使用多个线程和固定内存¶
使用torch.utils.data.DataLoader时，设置num_workers>0，而不是使用它的默认值0.并且设置pin_memory=True，而不是默认的False。具体的解释参见这里。
经验表明，当num_workers等于GPU数量的四倍时效果最好。注意这种方法会增大CPU负荷。
增大batch size¶
使用你的GPU所允许的最大的batch size可以提高你的训练速度。调整batch size时需要等比例的调整learning rate。
OpenAI对此有一篇很漂亮的论文可供参考。使用大batch size的缺点是，这可能会降低模型的泛化能力。
使用混合精度(AMP)¶
Pytorch 1.6增加了混合精度训练的官方实现。使用FP16和FP32混合精度可以训练地更快，而且相比单精度(FP32)训练并没有精度损失。下面是一个使用AMP的例子：
import torch
# Creates once at the beginning of training
scaler = torch.cuda.amp.GradScaler()
for data, label in data_iter:
optimizer.zero_grad()
# Casts operations to mixed precision
with torch.cuda.amp.autocast():
loss = model(data)
# Scales the loss, and calls backward()
# to create scaled gradients
scaler.scale(loss).backward()
# Unscales gradients and calls
# or skips optimizer.step()
scaler.step(optimizer)
# Updates the scale for next iteration
scaler.update()
AMP训练的benchmark可以参考这篇文章
考虑使用另一种优化器¶
AdamW往往比Adam有更好的优化效率。相关文章参见这篇。
此外还可以尝试最近异军突起的LAMB优化器。
如果使用英伟达的显卡，可以安装英伟达的APEX包，里面提供了针对NVIDA GPU而改良的常见优化器，比如Adam。
开启cudNN基准测试¶
如果你的模型架构是固定的，并且输入数据的size是常数，那么可以尝试设置torch.backends.cudnn.benchmark = True。这会启动cudNN autotuner来测试卷积计算的不同方法，并最终选用最佳方法。 需要注意的是，如果你把batch size最大化了，那么这个测试将会非常慢。
小心CPU和GPU之间频繁的数据转换¶
注意你的代码中是否频繁调用了tensor.cpu()或者tensor.cuda()，这会非常耗时。
如果你在创建一个新的tensor，你可以通过设置参数device=torch.device('cuda:0')来把数据直接创建在GPU上。
使用梯度累积¶
这种方法变相增加了batch size。如果你的GPU memory不足以容纳较大的batch，你可以把一个batch分几次输入，从而算出总梯度。你可以在调用optimizaer.step之前多次调用.backward，来实现这个方法。下面是一个例子。
model.zero_grad() # Reset gradients tensors
for i, (inputs, labels) in enumerate(training_set):
predictions = model(inputs) # Forward pass
loss = loss_function(predictions, labels) # Compute loss function
loss = loss / accumulation_steps # Normalize our loss (if averaged)
loss.backward() # Backward pass
if (i+1) % accumulation_steps == 0: # Wait for several backward steps
optimizer.step() # Now we can do an optimizer step
model.zero_grad() # Reset gradients tensors
if (i+1) % evaluation_steps == 0: # Evaluate the model when we...
evaluate_model() # ...have no gradients accumulated
使用Distributed Data Parallel 来进行多卡训练¶
使用torch.nn.DistributedDataParallel而不是torch.nn.DataParallel来开启多卡训练。这样每个GPU都会由一个独立的CPU核来驱动，避免了DataParallel的GIL问题。
将梯度设置为None而不是0¶
使用.zero_grad(set_to_none=True)而不是.zero_grad()。
这会让内存分配器处理梯度，而不是主动地将他们设置为0。这只会提供一个微量的加速，就像这篇文档所说的那样，所以不要期待有任何奇迹。
注意这么做会有一定副作用。请仔细查看文档。
使用.as_tensor()而不是.tensor()¶
torch.tensor()总是会复制数据。如果你想转化一个numpy数组，使用torch.as_tensor()或者torch.from_numpy()来避免拷贝数据。
使用梯度剪裁(gradient clipping)¶
这种方法一开始是用来避免RNN中的梯度爆炸，现在也有实验和理论说明梯度剪裁可以加速收敛。参见这篇文章。
抱抱脸团队实现的Transformer就是一个非常清晰的使用梯度裁剪的例子，这里面还用了本文提到的其他优化方法，比如AMP。
现在还不清楚这种方法适用于什么样的模型，不过目前来看它在RNN架构、基于Transformer的架构，以及ResNet架构上表现出很稳定的实用性。
在BatchNorm之前关掉bias¶
这非常简单：在BatchNormalization层之前关掉每一层的bias，也就是说,对于一个二维卷积层，把bias参数设为False：torch.nn.Conv2d(...,bias=False,...)。（原理参见这篇文档）
在做validation的时候关掉梯度计算¶
这很简单，在validation时设置torch.no_grad()
预加载数据¶
非常简单，用下面的DataLoaderX替换DataLoader即可：