Python tf.compat.v1.mixed_precision.enable_mixed_precision_graph_rewrite用法及代碼示例

通過圖形重寫啟用混合精度。

用法

tf.compat.v1.mixed_precision.enable_mixed_precision_graph_rewrite(
    opt, loss_scale='dynamic'
)

混合精度是在訓練模型以提高性能時同時使用 float32 和 float16 數據類型。這是通過圖形重寫操作和loss-scale 優化器實現的。

在 float16 中執行算術運算利用了專門的處理單元，例如 NVIDIA Tensor Cores，以獲得更高的算術吞吐量。但是，由於可表示範圍較小，使用 float16 執行整個訓練可能會導致梯度下溢，即小的梯度值變為零。相反，當使用兼容的硬件加速器時，僅在 float16 中執行選擇算術運算會導致更高的吞吐量和更少的訓練時間，同時還可以減少內存使用量，通常不會犧牲模型的準確性。

注意：雖然混合精度重寫改變了整個模型中各個層的數據類型，但預計在 float32 中達到的精度相同。如果在動態損失縮放中出現NaN 梯度，則跳過該批次的模型更新。在這種情況下，全局步數不會增加，LossScaleOptimizer 會嘗試減小損失縮放值以避免後續迭代中的 NaN 值。這種方法已被證明可以實現與 float32 相同的準確度，並且在大多數情況下，可以實現更好的訓練吞吐量。

例子：

model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(64, activation='softmax'),
])

opt = tf.keras.optimizers.SGD()
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
model.compile(loss="mse", optimizer=opt)

x_train = np.random.random((1024, 64))
y_train = np.random.random((1024, 64))
model.fit(x_train, y_train)

調用 enable_mixed_precision_graph_rewrite(opt) 在計算梯度之前啟用圖形重寫操作。該函數還返回一個用 LossScaleOptimizer 包的 Optimizer ( opt )。這可以防止在反向傳遞期間 float16 張量中的下溢。必須將 tf.train.Optimizer 或 tf.keras.optimizers.Optimizer 類型的優化器傳遞給此函數，然後將其包裝以使用損失縮放。

圖重寫操作將圖中某些操作的dtype從float32更改為float16。此重寫操作包含或排除了幾類操作。以下類別的 Ops 在auto_mixed_precision_lists.h 中類AutoMixedPrecisionLists 下的相應函數中定義：

• ClearList：在數值上沒有顯著不利影響的操作。例如： ArgMax 和 Floor 。
• AllowList ：被認為在 float16 中執行數字安全的操作，因此總是被轉換。例如： Conv2D 。
• DenyList ：在 float16 中執行的操作在數值上不安全，可能會對下遊節點產生負麵影響。例如： Softmax 。
• GrayList ：被認為在 float16 中執行數字安全的操作，除非來自 DenyList Op 的下遊。例如： Add 和 AvgPool 。

使用此函數時，隻能通過調用 opt.minimize() 或 opt.compute_gradients() 後跟 opt.apply_gradients() 來計算和應用梯度。不應使用 tf.gradients 或 tf.GradientTape 計算梯度。這是因為返回的優化器將應用損失縮放，而 tf.gradients 或 tf.GradientTape 不會。如果您直接使用 tf.gradients 或 tf.GradientTape ，您的模型可能會由於 float16 下溢問題而無法收斂。

啟用 Eager Execution 後，混合精度圖形重寫僅在 tf.function s 內啟用，因為在 tf.function s 之外，沒有圖形。

對於具有 Tensor 內核的 NVIDIA GPU，作為一般性能指南，如果小於 256，維度(例如批量大小、輸入大小、輸出大小和通道數)應該是 2 的冪，或者如果大於 256，則可以被 8 整除。信息，請查看 NVIDIA 深度學習性能指南。

目前，混合精度僅在具有計算能力 7.0 及更高版本(Volta、Turing 或更新的架構)的 NVIDIA Tensor Core GPU 上啟用。 CPU 和 TPU 上的圖形部分不會被圖形重寫所觸及。