使用 自動混合精度 加速訓練 tensorflow, pytorch

啓用方式請看官方介紹頁面：
https://developer.nvidia.com/automatic-mixed-precision

---

tensorflow 最爽，設置一個標誌就直接啓用了（是假的）。

2019/10/16更新
似乎 tf 1.13 和 更舊版本都不支持混合精度
建議使用 tf 1.14 及以上版本

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用法相似。不過我不會用這個框架，所以直接跳過。

---