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.