使用 混合精度 加速训练 tensorflow, pytorch, mxnet

启用方式请看官方介绍页面：
https://developer.nvidia.com/automatic-mixed-precision

---

tensorflow 最爽，设置一个标志就直接启用了。
跟据下面 pytorch 的实验结果，我估计对小模型来说效果差不多，所以没进行测试实验了。

2019/8/12更新：
做了个实验，发现有点坑。Volta以下显卡，需要设置环境变量 TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_IGNORE_PERFORMANCE=1
原因来自：https://devtalk.nvidia.com/default/topic/1052688/container-tensorflow/issue-about-no-suitable-gpus-detected-when-using-mixed-precision-graph-optimizer/
另外，直接设置环境变量 TF_ENABLE_AUTO_MIXED_PRECISION 好像并不会启动混合精度。这个坑。。。

启动混合精度训练后，观察控制台，可以看到以下类似输出

2019-08-12 22:00:20.533211: I tensorflow/core/grappler/optimizers/auto_mixed_precision.cc:1767] Running auto_mixed_precision graph optimizer
2019-08-12 22:00:20.537642: I tensorflow/core/grappler/optimizers/auto_mixed_precision.cc:1241] No whitelist ops found, nothing to do
2019-08-12 22:00:20.934960: I tensorflow/core/grappler/optimizers/auto_mixed_precision.cc:1767] Running auto_mixed_precision graph optimizer
2019-08-12 22:00:20.941479: I tensorflow/core/grappler/optimizers/auto_mixed_precision.cc:1723] Converted 25/262 nodes to float16 precision using 2 cast(s) to float16 (excluding Const and Variable casts)

使用混合精度训练的实验代码

import os
os.environ['TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_IGNORE_PERFORMANCE'] = '1'
import tensorflow as tf
from tensorflow import keras as K
import numpy as np

in_p = tf.placeholder(tf.float32, [None, 128, 128, 3])
net = K.layers.Conv2D(3, 32, 2, 'same', activation=K.activations.elu)(in_p)
net = K.layers.Conv2D(3, 64, 2, 'same', activation=K.activations.elu)(net)
net = K.layers.Conv2D(3, 128, 2, 'same', activation=K.activations.elu)(net)
net = K.layers.Conv2D(3, 256, 2, 'same', activation=K.activations.elu)(net)
net_out = tf.reduce_mean(net, axis=[1,2,3])

loss_op = tf.reduce_mean(tf.abs(net_out - tf.ones_like(net_out)))

optim = tf.train.AdamOptimizer()
optim = tf.train.experimental.enable_mixed_precision_graph_rewrite(optim)
optim_op = optim.minimize(loss_op)

s = tf.Session()
s.run(tf.global_variables_initializer())

for i in range(1000):
    x = np.random.normal(0, 1, [10, 128, 128, 3])
    loss, _ = s.run([loss_op, optim_op], {in_p: x})
    print(loss)

---

pytorch 这个需要一个apex库
github 链接 https://github.com/nvidia/apex
编译需要有 cuda 开发库 和 cpp编译器
不使用 cuda 开发库 和 cpp编译器 也可以编译，不过会缺失一些功能。
已测试。对于简单的模型可以很容易使用。不过遇到复杂的多模型的情况就会报一些奇怪错误。
我测试使用的是简单的堆叠模型做 cifar10 实验，模型也很小，没有观察到显存占用变小，不过观察到，貌似训练速度略微变慢了一点。。。这个可能是因为我的显卡本质上不支持半精度运算的原因。

---

mxnet，从介绍看也需要apex库。跟pytorch用法相似。不过我不会用这个框架，所以直接跳过。

---