Функция cuda.memory_stats

Функция cuda.memory_stats предназначена для получения детальной информации о распределении памяти CUDA. Она возвращает словарь, содержащий различные метрики, такие как количество выделенных блоков, размер активной памяти, пиковые значения и другие статистические данные. Основное применение функции - диагностика утечек памяти и оптимизация использования видеопамяти.

Синтаксис

torch.cuda.memory_stats(device=None)

Параметр device - необязательный, определяет устройство (по умолчанию используется текущее). Возвращает словарь с ключами-строками и значениями-числами.

Пример

Посмотрим базовую статистику для текущего устройства после выделения памяти:

import torch

torch.manual_seed(0)
if torch.cuda.is_available():
	stats = torch.cuda.memory_stats()
	print("Количество ключей в статистике:", len(stats))
	print("Пример ключа:", list(stats.keys())[0])
else:
	print("CUDA недоступна")

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

"Количество ключей в статистике: 47"
"Пример ключа: all.allocated_blocks"

Пример

Создадим тензор и измерим изменения в статистике выделенной памяти:

import torch

if torch.cuda.is_available():
	stats_before = torch.cuda.memory_stats()
	allocated_before = stats_before.get("allocated_bytes.all.current", 0)
	
	t = torch.randn(1000, 1000, device='cuda')
	
	stats_after = torch.cuda.memory_stats()
	allocated_after = stats_after.get("allocated_bytes.all.current", 0)
	
	print("До выделения:", allocated_before)
	print("После выделения:", allocated_after)
	print("Разница:", allocated_after - allocated_before)
else:
	print("CUDA недоступна")

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

"До выделения: 0"
"После выделения: 4000000"
"Разница: 4000000"

Выделение тензора размером 1000x1000 типа float32 занимает 4 000 000 байт (4 МБ).

Пример

Получим статистику для конкретного устройства и извлечём ключевые метрики:

import torch

if torch.cuda.is_available():
	stats = torch.cuda.memory_stats(device=0)
	keys = [
		"allocated_bytes.all.peak",
		"reserved_bytes.all.peak",
		"active_bytes.all.current"
	]
	for key in keys:
		if key in stats:
			print(f"{key}: {stats[key]}")
else:
	print("CUDA недоступна")

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

"allocated_bytes.all.peak: 4000000"
"reserved_bytes.all.peak: 4194304"
"active_bytes.all.current: 4000000"

Ключевые метрики: пик выделенной памяти, пик зарезервированной памяти и текущая активная память.

Пример

Используем статистику для мониторинга памяти внутри цикла обучения:

import torch

if torch.cuda.is_available():
	model = torch.nn.Linear(10, 10).cuda()
	data = torch.randn(100, 10, device='cuda')
	
	for step in range(3):
		output = model(data)
		loss = output.sum()
		loss.backward()
		
		stats = torch.cuda.memory_stats()
		allocated = stats.get("allocated_bytes.all.current", 0)
		reserved = stats.get("reserved_bytes.all.current", 0)
		
		print(f"Шаг {step}: выделено {allocated} байт, зарезервировано {reserved} байт")
		
		# Очистка градиентов для следующей итерации
		model.zero_grad()
else:
	print("CUDA недоступна")

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

"Шаг 0: выделено 2080 байт, зарезервировано 2097152 байт"
"Шаг 1: выделено 2080 байт, зарезервировано 2097152 байт"
"Шаг 2: выделено 2080 байт, зарезервировано 2097152 байт"

Статистика помогает отслеживать стабильность потребления памяти на протяжении обучения.

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

функцию memory_allocated,
которая возвращает текущий объём выделенной памяти
функцию memory_reserved,
которая возвращает объём зарезервированной памяти
функцию max_memory_allocated,
которая возвращает пик выделенной памяти
функцию empty_cache,
которая освобождает кэш памяти CUDA