Функция cuda.reset_peak_memory_stats

Функция cuda.reset_peak_memory_stats сбрасывает пиковые значения статистики памяти для указанного устройства CUDA. Она обнуляет значения, которые возвращают функции cuda.max_memory_allocated и cuda.max_memory_reserved. Параметр device указывает целевое устройство (по умолчанию - текущее устройство). Это позволяет изолированно измерять пиковое потребление памяти на определённых участках кода.

Синтаксис

torch.cuda.reset_peak_memory_stats(device=None)

Пример

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

import torch

torch.manual_seed(0)
device = 'cuda' if torch.cuda.is_available() else 'cpu'

torch.cuda.reset_peak_memory_stats()
print("Peak memory allocated after reset:", torch.cuda.max_memory_allocated())

t = torch.randn(1000, 1000).to(device)
print("Peak memory allocated after tensor creation:", torch.cuda.max_memory_allocated())

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

Peak memory allocated after reset: 0
Peak memory allocated after tensor creation: 4000000

Пример

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

import torch

torch.manual_seed(0)
device = 'cuda' if torch.cuda.is_available() else 'cpu'

t = torch.randn(500, 500).to(device)

torch.cuda.reset_peak_memory_stats()

for i in range(5):
	res = t * 2

print("Peak memory allocated after loop:", torch.cuda.max_memory_allocated())

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

Peak memory allocated after loop: 1000000

Пример

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

import torch

torch.manual_seed(0)

if torch.cuda.is_available():
	torch.cuda.set_device(0)
	t = torch.randn(2000, 2000).to('cuda:0')
	
	torch.cuda.reset_peak_memory_stats(device=0)
	print("Peak memory after reset on device 0:", torch.cuda.max_memory_allocated(0))

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

Peak memory after reset on device 0: 0

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

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