Tensorflow模型优化训练思路

问题现状

随着深度学习模型越来越大，数据集越来越大，模型的训练变得越来越慢。这对于想要快速验证算法的研究人员来说，是个比较麻烦的问题。

那么一般来说，我们会想要优化模型训练，以期更快验证模型效果。

无论是使用Tensorflow还是Pytorch来搭建模型，基本的训练优化思路都是一致的，只是由于框架的不同，某些优化细节有些差别。

这里探讨的是模型在GPU上的训练优化。

基本的训练优化思路

• Step 1 优化模型在单机单卡上的训练
• Step 2 优化模型在单机多卡上的训练
• Step 3 优化模型在多级多卡上的训练

Tensorflow模型的训练优化

针对基本思路，我们发掘一些优化细节。

Step 1 优化模型在单机单卡上的训练

首先，我们需要先将模型在单卡上的训练进行足够的优化，再去谈论扩展GPU数量来提升训练速度。

将数据预处理放在CPU上可以显著提高性能，这样可以让GPU专注训练，使用nvidia-smi来查看GPU的利用率是否达到80%~100%。

对于tensorflow来说，做到以下几点基本可以立马加速模型训练：

1. 使用Pinned memory
2. 打开AMP 
3. 打开XLA （input size需要是固定的）
4. 使用LAMB作为optmizer（而不用ADAM）
5. 使用TF32 
6. 在没有显存溢出的情况下，尝试更大的batch size
7. 融合op，以减少D2D/H2D/D2H的数据传输 （使用nsight system来profile模型ÿ