apex和pytorch autocast混合精度训练速度和时间对比

众所周知，自从pytorch更新1.6版本之后，就引入了自家的混合精度训练模式（下面统称autocast），这种方法简单，只需要添加几行代码即可实现，可以说是pytorch爱好者的福音。autocast在面对主流的apex时，是完爆还是被吊打呢，请看下面分析：

• 模型架构：Transformer-xl
• 环境：python3.6.9 、torch1.7.1、cuda10.1、显卡GeForce RTX 2080Ti

1、即不用apex，也不用autocast

• batch size = 1
  
  占用内存：7561M，每个batch的时间为285.69ms

• batch size = 2
  

占用内存：10323M，每个batch的时间为503.74ms

2、autocast

• batch size = 1
  

占用内存：7515M，每个batch的时间为272.65ms

• batch size = 2
  

占用内存：8965M，每个batch的时间为451.57ms

3、apex

• batch size = 1
  
  占用内存：7831M，每个batch的时间为246.25ms

• batch size = 2
  
  占用内存：9521M，每个batch的时间为389.40ms

4、结果

• 空间上：在普通训练时，增加一个batch，会增加10323-7561=2762M内存，autocast增加1450M，apex增加1690M。相当于：普通训练 = 1.9 * autocast = 1.6 * apex

• 时间上：bsz = 2时，普通训练 = 1.12 * autocast = 1.29 * apex；
  bsz = 1时，普通训练 = 1.05 * autocast = 1.16 * apex；因为机器内存有限，这里只能测试batch size为1和2的情况，但是也可以发现问题，当batch size变大时，普通训练和autocast，apex的差别变得更大了，也就是说，autocast和apex在batch size越大的时候，优势越明显。

• autocast相对于apex更省内存，但是速度上会稍微慢一些。

	方式	普通训练	autocast	apex
bsz = 1	内存(M)	7561	7515	7831
bsz = 2	内存(M)	10323	8965	9521
bsz = 1	时间(ms)	285.69	272.65	246.25
bsz = 2	时间(ms)	503.74	451.57	389.40

补充

之后，我在3090显卡上测试，使用apex可以放9个batch，用autocast可以放11个batch。时间上，autocast是1.82s/batch，apex是1.75s/batch，说明，在batch比较大的时候，autocast相比apex在显存上是有比较大的优势的，时间上都相差不大。因此，鉴于apex的安装步骤繁琐困难，建议大家使用autocast进行训练。