Функция cuda.max_memory_allocated

Функция cuda.max_memory_allocated возвращает максимальный (пиковый) объём памяти GPU в байтах, которая была выделена для тензоров на текущем устройстве с момента последнего вызова cuda.reset_peak_memory_stats или с момента начала работы программы. В отличие от cuda.memory_allocated, которая показывает текущее потребление, эта функция позволяет оценить пиковую нагрузку на видеопамять, что полезно для оптимизации размера батча или архитектуры модели. Функция не принимает аргументов и возвращает целое число типа int.

Синтаксис

torch.cuda.max_memory_allocated([device])

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

Пример

Давайте создадим несколько тензоров на GPU и посмотрим на изменение пикового значения выделенной памяти:

import torch

# Убедимся, что CUDA доступна
if torch.cuda.is_available():
	# Сбросим пиковые значения для чистоты эксперимента
	torch.cuda.reset_peak_memory_stats()
	
	# Создаём тензор на GPU
	t1 = torch.randn(1000, 1000).cuda()
	print(f"Пиковое выделение: {torch.cuda.max_memory_allocated()} байт")
	
	# Создаём ещё один тензор
	t2 = torch.randn(2000, 2000).cuda()
	print(f"Пиковое выделение после t2: {torch.cuda.max_memory_allocated()} байт")
	
	# Удаляем тензоры
	del t1, t2
	torch.cuda.empty_cache()
	
	# Пиковое значение не изменилось
	print(f"Пиковое выделение после удаления: {torch.cuda.max_memory_allocated()} байт")

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

Пиковое выделение: 4000000 байт
Пиковое выделение после t2: 20000000 байт
Пиковое выделение после удаления: 20000000 байт

Обратите внимание, что пиковое значение остаётся на максимальном достигнутом уровне, даже после освобождения памяти.

Пример

Рассмотрим, как сбросить пиковые значения с помощью cuda.reset_peak_memory_stats и получить новые показатели:

import torch

if torch.cuda.is_available():
	# Создаём тензор
	t = torch.randn(500, 500).cuda()
	print(f"Пик после создания: {torch.cuda.max_memory_allocated()} байт")
	
	# Сбрасываем пиковую статистику
	torch.cuda.reset_peak_memory_stats()
	
	# Создаём новый тензор
	t2 = torch.randn(800, 800).cuda()
	print(f"Новый пик после сброса: {torch.cuda.max_memory_allocated()} байт")

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

Пик после создания: 1000000 байт
Новый пик после сброса: 2560000 байт

Как видно, после сброса статистика начинает отсчёт заново, и пиковое значение отражает только новые выделения памяти.

Пример

Использование функции с указанием конкретного устройства GPU (если их несколько):

import torch

if torch.cuda.is_available():
	device_count = torch.cuda.device_count()
	for i in range(device_count):
		# Переключаемся на устройство i
		torch.cuda.set_device(i)
		# Создаём тензор
		t = torch.randn(100, 100).cuda(i)
		# Получаем пиковое выделение для этого устройства
		peak = torch.cuda.max_memory_allocated(i)
		print(f"Устройство {i}: пиковое выделение = {peak} байт")

Результат выполнения кода (для системы с двумя GPU):

Устройство 0: пиковое выделение = 40000 байт
Устройство 1: пиковое выделение = 40000 байт

Это позволяет отслеживать пиковую память на каждом устройстве независимо.

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

функцию cuda.memory_allocated,
которая возвращает текущий объём занятой памяти GPU
функцию cuda.memory_reserved,
которая показывает объём зарезервированной кеш-памяти
функцию cuda.reset_peak_memory_stats,
которая сбрасывает пиковые значения статистики памяти
функцию cuda.empty_cache,
которая освобождает неиспользуемую кеш-память GPU