Функция cuda.is_tf32_supported

Функция cuda.is_tf32_supported проверяет, поддерживает ли текущее устройство CUDA формат TF32 (TensorFloat-32). TF32 - это математический формат, представленный NVIDIA в архитектуре Ampere, который сочетает 19-битную мантиссу от FP32 и 8-битную экспоненту от FP16. Функция не принимает параметров и возвращает булево значение.

Синтаксис

torch.cuda.is_tf32_supported()

Возвращаемое значение

Функция возвращает bool - True, если устройство поддерживает TF32, и False в противном случае.

Пример

Давайте проверим, поддерживает ли текущее CUDA-устройство TF32:

import torch

is_supported = torch.cuda.is_tf32_supported()
print(is_supported)

Результат выполнения кода (зависит от аппаратного обеспечения):

True

Пример

Перед выполнением операций с TF32 можно проверить поддержку и установить точность вычислений для матричных операций:

import torch

if torch.cuda.is_tf32_supported():
	torch.backends.cuda.matmul.allow_tf32 = True
	torch.backends.cudnn.allow_tf32 = True
	print("TF32 enabled for matmul and cudnn")
else:
	print("TF32 not supported on this device")

Результат выполнения кода:

"TF32 enabled for matmul and cudnn"

Пример

Проверим поддержку TF32 на всех доступных устройствах CUDA:

import torch

num_devices = torch.cuda.device_count()
for i in range(num_devices):
	torch.cuda.set_device(i)
	supported = torch.cuda.is_tf32_supported()
	name = torch.cuda.get_device_name(i)
	print(f"Device {i}: {name} - TF32 supported: {supported}")

Результат выполнения кода:

"Device 0: NVIDIA A100-SXM4-40GB - TF32 supported: True"

Примечания

Функция is_tf32_supported возвращает True только для устройств с архитектурой NVIDIA Ampere (A100) или новее. TF32 значительно ускоряет вычисления на таких устройствах, но может незначительно снижать точность. Включение TF32 регулируется флагами torch.backends.cuda.matmul.allow_tf32 и torch.backends.cudnn.allow_tf32.

Смотрите также

функцию is_available,
которая проверяет доступность CUDA в системе
функцию device_count,
которая возвращает количество доступных CUDA-устройств
функцию get_device_name,
которая возвращает название текущего CUDA-устройства
функцию is_bf16_supported,
которая проверяет поддержку формата BF16 на устройстве