Функция cuda.is_bf16_supported

Функция cuda.is_bf16_supported проверяет, поддерживает ли текущее CUDA-устройство формат BF16 (bfloat16). Она не принимает параметров и возвращает логическое значение True или False. Формат BF16 используется для ускорения вычислений и экономии памяти в глубоком обучении.

Синтаксис

torch.cuda.is_bf16_supported()

Пример

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

import torch

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

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

True

Пример

Проверим поддержку BF16 перед созданием тензора с этим типом данных:

import torch

if torch.cuda.is_bf16_supported():
	t = torch.tensor([1, 2, 3, 4, 5], dtype=torch.bfloat16)
	print(t)
else:
	print("BF16 not supported")

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

tensor([1., 2., 3., 4., 5.], dtype=torch.bfloat16)

Пример

Проверим поддержку BF16 с использованием конкретного устройства:

import torch

device = torch.device("cuda:0")
if torch.cuda.is_bf16_supported():
	t = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.bfloat16)
	t = t.to(device)
	print(t.device)
else:
	print("BF16 not supported")

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

cuda:0

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

функцию is_available,
которая проверяет доступность CUDA
функцию is_tf32_supported,
которая проверяет поддержку TF32
функцию get_device_name,
которая возвращает имя устройства
функцию device_count,
которая возвращает количество устройств