Функция cuda.max_memory_reserved

Функция cuda.max_memory_reserved возвращает максимальный объём памяти на GPU, который был зарезервирован кеширующим аллокатором PyTorch с момента начала выполнения программы или с момента последнего вызова функции cuda.reset_peak_memory_stats. Зарезервированная память - это память, которую аллокатор выделил у драйвера CUDA и удерживает в своём кеше для последующих выделений, чтобы избежать затратных обращений к драйверу. Эта функция не принимает никаких параметров и возвращает целое число - количество байт.

Синтаксис

torch.cuda.max_memory_reserved([device])

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

Пример

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

import torch

torch.manual_seed(0)
if torch.cuda.is_available():
	t = torch.randn(1000, 1000, device='cuda')
	peak_reserved = torch.cuda.max_memory_reserved()
	print(peak_reserved)

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

4194304

Пример

Покажем, как изменяется значение при последовательных выделениях памяти:

import torch

torch.manual_seed(0)
if torch.cuda.is_available():
	t1 = torch.randn(500, 500, device='cuda')
	print(torch.cuda.max_memory_reserved())
	
	t2 = torch.randn(1000, 1000, device='cuda')
	print(torch.cuda.max_memory_reserved())

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

2097152
4194304

Пример

Используем функцию для конкретного устройства GPU с индексом 0:

import torch

torch.manual_seed(0)
if torch.cuda.is_available():
	t = torch.randn(2000, 2000, device='cuda:0')
	peak_reserved = torch.cuda.max_memory_reserved(0)
	print(peak_reserved)

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

16777216

Пример

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

import torch

torch.manual_seed(0)
if torch.cuda.is_available():
	t1 = torch.randn(1000, 1000, device='cuda')
	print(torch.cuda.max_memory_reserved())
	
	torch.cuda.reset_peak_memory_stats()
	t2 = torch.randn(500, 500, device='cuda')
	print(torch.cuda.max_memory_reserved())

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

4194304
2097152

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

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