解决Pytorch半精度浮点型网络训练的问题

用pytorch1.0进行半精度浮点型网络训练需要注意下问题：

1、网络要在gpu上跑，模型和输入样本数据都要cuda().half()

2、模型参数转换为half型，不必索引到每层，直接model.cuda().half()即可

3、对于半精度模型，优化算法，adam我在使用过程中，在某些参数的梯度为0的时候，更新权重后，梯度为零的权重变成了nan，这非常奇怪，但是adam算法对于全精度数据类型却没有这个问题。

另外，sgd算法对于半精度和全精度计算均没有问题。

还有一个问题是不知道是不是网络结构比较小的原因，使用半精度的训练速度还没有全精度快。这个值得后续进一步探索。

对于上面的这个问题，的确是网络很小的情况下，在1080ti上半精度浮点型没有很明显的优势，但是当网络变大之后，半精度浮点型要比全精度浮点型要快。

但具体快多少和模型的大小以及输入样本大小有关系，我测试的是要快1/6，同时，半精度浮点型在占用内存上比较有优势，对于精度的影响尚未探究。

将网络再变大些，epoch的次数也增大，半精度和全精度的时间差就表现出来了，在训练的时候。

补充：pytorch半精度，混合精度，单精度训练的区别amp.initialize

看代码吧~

为了帮助提高pytorch的训练效率，英伟达提供了混合精度训练工具apex。号称能够在不降低性能的情况下，将模型训练的速度提升2-4倍，训练显存消耗减少为之前的一半。

文档地址是：

该 工具 提供了三个功能，amp、parallel和normalization。由于目前该工具还是0.1版本，功能还是很基础的，在最后一个normalization功能中只提供了layernorm层的复现，实际上在后续的使用过程中会发现，出现问题最多的是pytorch的bn层。

第二个工具是pytorch的分布式训练的复现，在文档中描述的是和pytorch中的实现等价，在代码中可以选择任意一个使用，实际使用过程中发现，在使用混合精度训练时，使用apex复现的parallel工具，能避免一些bug。

默认训练方式是 单精度float32

半精度 model(img.half())

接下来是混合精度的实现，这里主要用到apex的amp工具。

代码修改为：

加上这一句封装，

实际流程为：调用amp.initialize按照预定的opt_level对model和optimizer进行设置。在计算loss时使用amp.scale_loss进行回传。

需要注意以下几点：

在调用amp.initialize之前，模型需要放在gpu上，也就是需要调用cuda()或者to()。

在调用amp.initialize之前，模型不能调用任何分布式设置函数。

此时输入数据不需要在转换为半精度。

在使用混合精度进行计算时，最关键的参数是opt_level。他一共含有四种设置值：‘00'，‘01'，‘02'，‘03'。实际上整个amp.initialize的输入参数很多：

但是在实际使用过程中发现，设置opt_level即可，这也是文档中例子的使用方法，甚至在不同的opt_level设置条件下，其他的参数会变成无效。（已知bug：使用‘01'时设置keep_batchnorm_fp32的值会报错）

概括起来：

00相当于原始的单精度训练。01在大部分计算时采用半精度，但是所有的模型参数依然保持单精度，对于少数单精度较好的计算（如softmax）依然保持单精度。02相比于01，将模型参数也变为半精度。

03基本等于最开始实验的全半精度的运算。值得一提的是，不论在优化过程中，模型是否采用半精度，保存下来的模型均为单精度模型，能够保证模型在其他应用中的正常使用。这也是apex的一大卖点。

在pytorch中，bn层分为train和eval两种操作。

实现时若为单精度网络，会调用cudnn进行计算加速。常规训练过程中bn层会被设为train。apex优化了这种情况，通过设置keep_batchnorm_fp32参数，能够保证此时bn层使用cudnn进行计算，达到最好的计算速度。

但是在一些fine tunning场景下，bn层会被设为eval（我的模型就是这种情况）。此时keep_batchnorm_fp32的设置并不起作用，训练会产生数据类型不正确的bug。此时需要人为的将所有bn层设置为半精度，这样将不能使用cudnn加速。

一个设置的参考代码如下：

实际测试下来，最后的模型准确度上感觉差别不大，可能有轻微下降；时间上变化不大，这可能会因不同的模型有差别；显存开销上确实有很大的降低。

以上为个人经验，希望能给大家一个参考，也希望大家多多支持。