理論
在混合精度訓練中，權重，激活值和梯度是保存成fp16的形式，為了能夠匹配fp32的網絡精度，有一個權重的fp32的master copy。
在tensorflow中的具體實現
tensorflow支持fp16的存儲和tensor計算。包含tf.float16的數據類型的卷積和矩陣運算會自動使用fp16的計算。
為了能夠使用tensor的core，fp32的模型需要轉換成fp32和fp16的混合，可以手動完成，也可以自動混合精度(AMP)。
Tensorflow中自動混合精度訓練
自動混合精度訓練的啟動很簡單，只需要設定一個環境變量 export tf_enable_auto_mixed_precision=1
當然也可以在tensorflow的python腳本里面去設定 os.environ['TF_ENABLE_AUTO_MIXED_PRECISION']='1'
一旦打開之后，AMP就會幫你做以下兩件事情
1.在你的tensorflow的圖中插入合適的cast運算，在合適的時機使用float16來計算和存儲
2.在訓練optimizer里面，打開自動loss scaling
手動在tensorflow中使用混合精度訓練
1.去nvidia GPU cloud(NGC) container registry中拿最新的tensorflow container，這個container已經build好了，測試過並且調試過，可以直接用。這個contianer包含最新的cuda版本，fp16支持，並且基於最新的架構進行了優化
2.在矩陣運算或者卷積運算中采用tf.fp16，這個數據類型會盡可能的使用tensor core的硬件，舉個例子
dtype = tf.float16
data = tf.placeholder(dtype, shape=(nbatch, nin))
weights = tf.get_variable('weights', (nin, nout), dtype)
biases = tf.get_variable('biases', nout, dtype,
initializer=tf.zeros_initializer())
logits = tf.matmul(data, weights) + biases
3.確保可以訓練的參數是fp32類型的，並且在模型中使用他們的時候cast成fp16
tf.cast(tf.get_variable(..., dtype=tf.float32), tf.float16)
4.確保softmax的計算時fp32精度的
tf.losses.softmax_cross_entropy(target, tf.cast(logits, tf.float32))
5.使用loss-scaling，這個會在計算梯度的時候乘以一個因子，然后算出來之后除以一個相同的因子
loss, params = ... scale = 128 grads = [grad / scale for grad in tf.gradients(loss * scale, params)]
參考：
https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html#tensorflow
https://arxiv.org/pdf/1710.03740.pdf