Флаг backends.cudnn.benchmark

Флаг torch.backends.cudnn.benchmark управляет автоматическим выбором наиболее быстрых алгоритмов для сверточных операций в библиотеке CuDNN. Если установить значение True, то PyTorch будет запускать несколько алгоритмов на этапе прогрева и выбирать самый производительный для конкретной конфигурации тензоров. Это позволяет ускорить обучение моделей, но требует, чтобы размеры входных тензоров оставались постоянными. При изменении размеров флаг пересчитывает оптимальные алгоритмы, что может вызывать задержки. Данный флаг имеет смысл включать для моделей с фиксированными размерами батчей и входных данных.

Синтаксис

import torch

# Включение автоматической оптимизации
torch.backends.cudnn.benchmark = True

# Отключение автоматической оптимизации
torch.backends.cudnn.benchmark = False

# Проверка текущего состояния флага
status = torch.backends.cudnn.benchmark
print(status)

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

False

Пример

Давайте рассмотрим простейший пример включения и проверки состояния флага:

import torch

# Проверяем текущее состояние
print("Initial:", torch.backends.cudnn.benchmark)

# Включаем оптимизацию
torch.backends.cudnn.benchmark = True
print("After set:", torch.backends.cudnn.benchmark)

# Отключаем оптимизацию
torch.backends.cudnn.benchmark = False
print("After reset:", torch.backends.cudnn.benchmark)

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

Initial: False
After set: True
After reset: False

Пример

Рассмотрим применение флага на реальной модели с сверточными слоями:

import torch
import torch.nn as nn

# Включаем оптимизацию CuDNN
torch.backends.cudnn.benchmark = True

class ConvNet(nn.Module):
	def __init__(self):
		super(ConvNet, self).__init__()
		self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
		self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
		self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
	
	def forward(self, x):
		x = torch.relu(self.conv1(x))
		x = torch.relu(self.conv2(x))
		x = torch.relu(self.conv3(x))
		return x

# Создаём модель и фиксированный вход
model = ConvNet().cuda()
t = torch.randn(8, 3, 224, 224).cuda()

# Выполняем несколько прямых проходов для прогрева
for _ in range(5):
	res = model(t)

print("Training with benchmark enabled")
print("Output shape:", res.shape)

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

Training with benchmark enabled
Output shape: torch.Size([8, 64, 224, 224])

Пример

Покажем, как изменение размеров батча влияет на работу флага:

import torch
import torch.nn as nn
import time

torch.backends.cudnn.benchmark = True

# Простая сверточная модель
model = nn.Conv2d(3, 64, kernel_size=3, padding=1).cuda()

# Разные размеры батча
shapes = [(4, 3, 128, 128), (8, 3, 128, 128), (16, 3, 128, 128)]

for shape in shapes:
	t = torch.randn(shape).cuda()
	
	# Прогрев
	for _ in range(10):
		res = model(t)
	
	# Замер времени
	start = time.time()
	for _ in range(100):
		res = model(t)
	elapsed = time.time() - start
	
	print(f"Batch size {shape[0]}: {elapsed:.4f} sec")

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

Batch size 4: 0.0023 sec
Batch size 8: 0.0034 sec
Batch size 16: 0.0051 sec

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

флаг cudnn.deterministic,
который обеспечивает воспроизводимость результатов
флаг cudnn.allow_tf32,
который включает использование TensorFloat-32 в CuDNN
функцию cudnn.is_available,
которая проверяет доступность библиотеки CuDNN
функцию cudnn.version,
которая возвращает версию установленной библиотеки CuDNN