tf.config.experimental.enable_tensor_float_32_execution

Включить или выключить использование TensorFloat-32 на поддерживаемом оборудовании.

tf.config.experimental.enable_tensor_float_32_execution(
    enabled
)

TensorFloat-32, или TF32, — это математический режим для графических процессоров NVIDIA Ampere. Использование TensorFloat-32 приводит к значительному ускорению выполнения определенных операций float32, таких как умножение матриц и свертки, на графических процессорах Ampere, но с уменьшенной точностью. Практически это уменьшение точности не должно влиять на сходимость моделей глубокого обучения.

TensorFloat-32 включен по умолчанию. TensorFloat-32 поддерживается только на графических процессорах Ampere, поэтому на другом оборудовании будет использована полная точность float32 независимо от того, включено ли TensorFloat-32 или нет. Если вы хотите использовать полную точность float32 на Ampere, вы можете отключить выполнение 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 битов точности. Это позволяет операциям выполняться намного быстрее, используя тензорные ядра графического процессора. TensorFloat-32 имеет тот же динамический диапазон, что и float32, что означает, что он не более склонен к потере или переполнению, чем float32. При включенном TensorFloat-32 операции по-прежнему используют накопление float32. Включение или отключение TensorFloat-32 влияет только на графические процессоры Ampere и последующие графические процессоры, которые поддерживают TensorFloat-32.

Обратите внимание, что TensorFloat-32 не всегда используется в поддерживаемых операциях, так как поддерживаются только входные данные определённых форм. Поддержка дополнительных форм входных данных и дополнительных операций может быть добавлена в будущем. В результате точность операций float32 может снизиться в небольших версиях TensorFlow.

TensorFloat-32 также используется для некоторых операций complex64. В настоящее время TensorFloat-32 используется реже для операций complex64, чем для операций float32.