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看起来像这样：

1cycle

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即可：

from torch.utils.data import DataLoader
from prefetch_generator import BackgroundGenerator
class DataLoaderX(DataLoader):
    def __iter__(self):
        return BackgroundGenerator(super().__iter__())