用於深度學習的自動混合精度
深度神經網絡訓練傳統上依賴IEEE單精度格式,但在混合精度的情況下,可以訓練半精度,同時保持單精度網絡的精度。這種同時使用單精度和半精度表示的技術稱為混合精度技術。
混合精度訓練的好處
通過使用Tensor Core加速數學密集型運算,如線性和卷積層。
與單精度相比,通過訪問一半的字節可以加快內存受限的操作。
減少訓練模型的內存需求,支持更大的模型或更小的批。
啟用混合精度涉及兩個步驟:在適當的情況下,將模型移植到使用半精度數據類型;並使用損失縮放來保持較小的梯度值。
TensorFlow、PyTorch和MXNet中的自動混合精度特性為深度學習研究人員和工程師提供了在NVIDIA Volta和Turing gpu上最多3倍的人工智能訓練速度,而只需要添加幾行代碼。

使用自動混合精度的主要深度學習框架
- TensorFlow
在NVIDIA NGC容器注冊表中提供的TensorFlow容器中提供了自動混合精度特性。要在容器內啟用此功能,只需設置一個環境變量:
export TF_ENABLE_AUTO_MIXED_PRECISION=1
另外,環境變量可以在TensorFlow Python腳本中設置:
os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'
另外還需要對優化器(Optimizer)作如下修改:
optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite(optimizer) # 需要添加這句話,該例子是tf1.14.0版本,不同版本可能不一樣
自動混合精度在TensorFlow內部應用這兩個步驟,使用一個環境變量,並在必要時進行更細粒度的控制。
- PyTorch
自動混合精度特性在GitHub上的Apex repository中可用。要啟用,請將這兩行代碼添加到您現有的訓練腳本中:
model, optimizer = amp.initialize(model, optimizer)
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
- MXNet
NVIDIA正在為MXNet構建自動混合精度特性。你可以在GitHub上找到正在進行的工作。要啟用該功能,請在現有的訓練腳本中添加以下代碼行:
amp.init()
amp.init_trainer(trainer)
with amp.scale_loss(loss, trainer) as scaled_loss:
autograd.backward(scaled_loss)
