Флаг backends.cudnn.allow_tf32

Флаг torch.backends.cudnn.allow_tf32 управляет использованием формата TF32 (TensorFloat-32) в библиотеке cuDNN. Этот формат позволяет ускорить вычисления на графических процессорах NVIDIA с архитектурой Ampere и новее за счёт использования 19-битной мантиссы вместо 23-битной. По умолчанию флаг включён, что обеспечивает максимальную производительность, но может привести к незначительной потере точности.

Синтаксис

torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = False

Пример

Давайте проверим текущее состояние флага:

import torch

res = torch.backends.cudnn.allow_tf32
print(res)

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

True

Пример

Давайте отключим использование TF32 для обеспечения максимальной точности:

import torch

torch.backends.cudnn.allow_tf32 = False
res = torch.backends.cudnn.allow_tf32
print(res)

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

False

Пример

Включим использование TF32 для повышения производительности вычислений:

import torch

torch.backends.cudnn.allow_tf32 = True
res = torch.backends.cudnn.allow_tf32
print(res)

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

True

Пример

Давайте создадим свёрточный слой и выполним вычисления при включенном TF32:

import torch

torch.manual_seed(0)
torch.backends.cudnn.allow_tf32 = True
conv = torch.nn.Conv2d(3, 64, kernel_size=3)
t = torch.randn(1, 3, 32, 32)
res = conv(t)
print(res.shape)

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

torch.Size([1, 64, 30, 30])

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

флаг set_float32_matmul_precision,
который управляет точностью умножения матриц
флаг cudnn.benchmark,
который включает автоматический поиск оптимальных алгоритмов
флаг cudnn.deterministic,
который обеспечивает детерминированность вычислений
функцию is_tf32_supported,
которая проверяет поддержку TF32 на устройстве