pytorch模型训练加速tricks

1、学习率设置策略

Pytorch 已经实现了两种方法：「torch.optim.lr_scheduler.CyclicLR」和「torch.optim.lr_scheduler.OneCycleLR」。
参考文档：https://pytorch.org/docs/stable/optim.html

2、dataloader中使用多个worker和页锁定内存

当使用 torch.utils.data.DataLoader 时 num_workers > 0 pin_memory=True

worker 数量的经验法则是将其设置为可用 GPU 数量的四倍，大于或小于这个数都会降低训练速度。请注意，增加 num_workers 将增加 CPU 内存消耗。

参考文档：https://pytorch.org/docs/stable/data.html

3、增大batch

 batch 调到最大是一个颇有争议的观点。一般来说，如果在 GPU 内存允许的范围内将 batch 调到最大，你的训练速度会更快。

但是，你也必须调整其他超参数，比如学习率。一个比较好用的经验是，batch 大小加倍时，学习率也要加倍。

4、换一种优化器

AdamW 是由 fast.ai 推广的一种具有权重衰减（而不是 L2 正则化）的 Adam，在 PyTorch 中以 torch.optim.AdamW 实现。AdamW 似乎在误差和训练时间上都一直优于 Adam。
Adam 和 AdamW 都能与上面提到的 1Cycle 策略很好地搭配。

5、cudnn基准

如果你的模型架构保持不变、输入大小保持不变，设置 torch.backends.cudnn.benchmark = True

6、小心CPU和GPU之间数据的传输

当频繁地使用 tensor.cpu() 将张量从 GPU 转到 CPU（或使用 tensor.cuda() 将张量从 CPU 转到 GPU）时，代价是非常昂贵的。item() 和 .numpy() 也是一样可以使用. detach() 代替。
如果你创建了一个新的张量，可以使用关键字参数 device=torch.device('cuda:0') 将其分配给 GPU。
如果你需要传输数据，可以使用. to(non_blocking=True)，只要在传输之后没有同步点。

7、使用梯度积累

增加 batch 大小的另一种方法是在调用 optimizer.step() 之前在多个. backward() 传递中累积梯度。

Hugging Face 的 Thomas Wolf 的文章《Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups》介绍了如何使用梯度累积。梯度累积可以通过如下方式实现：

model.zero_grad()                                   
# Reset gradients tensorsfor 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 accumulate

这个方法主要是为了规避 GPU 内存的限制而开发的。

8、分布式数据并行进行多GPU训练

加速分布式训练可能有很多方法，但是简单的方法是使用 torch.nn.DistributedDataParallel 而不是 torch.nn.DataParallel。

这样一来，每个 GPU 将由一个专用的 CPU 核心驱动，避免了 DataParallel 的 GIL 问题。

分布式训练文档地址：https://pytorch.org/tutorials/beginner/dist_overview.html

9、设置梯度为None，而不是0

梯度设置为. zero_grad(set_to_none=True) 而不是 .zero_grad()。

这样做可以让内存分配器处理梯度，而不是将它们设置为 0。

正如文档中所说，将梯度设置为 None 会产生适度的加速，但不要期待奇迹出现。

注意，这样做也有缺点，详细信息请查看文档。

文档地址：https://pytorch.org/docs/stable/optim.html

10、使用.as_tensor() 而不是.tensor()

torch.tensor() 总是会复制数据。如果你要转换一个 numpy 数组，使用 torch.as_tensor() 或 torch.from_numpy() 来避免复制数据。

11、验证期间关闭梯度计算

设置：torch.no_grad()

12、使用输入和batch归一化