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。