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

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

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

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