Python tf.config.experimental.enable_tensor_float_32_execution用法及代碼示例

在支持的硬件上啟用或禁用 TensorFloat-32。

用法

tf.config.experimental.enable_tensor_float_32_execution(
    enabled
)

TensorFloat-32，簡稱 TF32，是 NVIDIA Ampere GPU 的數學模式。 TensorFloat-32 執行會導致某些 float32 操作(例如矩陣乘法和卷積)在 Ampere GPU 上運行得更快，但精度會降低。這種降低的精度不應影響深度學習模型在實踐中的收斂。

默認情況下啟用 TensorFloat-32。 TensorFloat-32 僅在 Ampere GPU 上受支持，因此無論是否啟用 TensorFloat-32，所有其他硬件都將使用完整的 float32 精度。如果要在 Ampere 上使用完整的 float32 精度，可以使用此函數禁用 TensorFloat-32 執行。例如：

x = tf.fill((2, 2), 1.0001)
y = tf.fill((2, 2), 1.)
# TensorFloat-32 is enabled, so matmul is run with reduced precision
print(tf.linalg.matmul(x, y))  # [[2., 2.], [2., 2.]]
tf.config.experimental.enable_tensor_float_32_execution(False)
# Matmul is run with full precision
print(tf.linalg.matmul(x, y))  # [[2.0002, 2.0002], [2.0002, 2.0002]]

要檢查當前是否啟用了 TensorFloat-32 執行，請使用tf.config.experimental.tensor_float_32_execution_enabled.

如果啟用了 TensorFloat-32，則在大多數情況下，支持的操作的 float32 輸入(例如 tf.linalg.matmul )將從 23 位精度四舍五入到 10 位精度。這允許操作通過利用 GPU 的張量核心更快地執行。 TensorFloat-32 與 float32 具有相同的動態範圍，這意味著它比 float32 更不可能下溢或上溢。啟用 TensorFloat-32 後，Ops 仍然使用 float32 累積。啟用或禁用 TensorFloat-32 僅影響 Ampere GPU 和支持 TensorFloat-32 的後續 GPU。

注意 TensorFloat-32 並不總是用於支持的操作，因為僅支持某些形狀的輸入。將來可能會添加對更多輸入形狀和更多操作的支持。因此，在 TensorFlow 的次要版本中，float32 操作的精度可能會降低。

TensorFloat-32 也用於一些複雜的 64 位運算。目前，TensorFloat-32 用於 complex64 的情況較少，因為它用於 float32。