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, таких как матричные умножения и свёртки, на таких графических процессорах, но с уменьшенной точностью. Как правило, это уменьшение точности не сказывается на сходимости моделей глубокого обучения.

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

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

Чтобы проверить, включено ли выполнение TensorFloat-32, используйте tf.config.experimental.tensor_float_32_execution_enabled.

Если TensorFloat-32 включен, входные float32 поддерживаемых операций, таких как tf.linalg.matmul, в большинстве случаев будут округляться с 23 бит точности до 10 бит точности. Это позволяет операциям выполняться намного быстрее, используя тензорные ядра графического процессора. TensorFloat-32 имеет такой же динамический диапазон, как float32, что означает, что он не более подвержен подпотоку или переполнению, чем float32. Операции по-прежнему используют накопление float32, когда включён TensorFloat-32. Включение или отключение TensorFloat-32 влияет только на графические процессоры Ampere и новее.

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

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

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