Функция cuda.mem_get_info

Функция cuda.mem_get_info возвращает кортеж из двух значений: количество свободной памяти и общий объем памяти на указанном устройстве CUDA. Первым параметром функция принимает идентификатор устройства (по умолчанию текущее устройство). Функция полезна для мониторинга использования памяти GPU и оптимизации размера батчей.

Синтаксис

torch.cuda.mem_get_info([device])

Функция возвращает кортеж (free, total), где free - свободная память в байтах, а total - общий объем памяти устройства в байтах.

Пример

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

import torch

if torch.cuda.is_available():
	free, total = torch.cuda.mem_get_info()
	print(f"Free memory: {free / 1024**3:.2f} GB")
	print(f"Total memory: {total / 1024**3:.2f} GB")

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

Free memory: 11.72 GB
Total memory: 15.78 GB

Пример

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

import torch

if torch.cuda.is_available():
	device_count = torch.cuda.device_count()
	if device_count > 1:
		free, total = torch.cuda.mem_get_info(1)
		print(f"Device 1 free: {free / 1024**3:.2f} GB")
		print(f"Device 1 total: {total / 1024**3:.2f} GB")

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

Device 1 free: 10.45 GB
Device 1 total: 15.78 GB

Пример

Давайте создадим тензор и проверим, как изменится свободная память:

import torch

if torch.cuda.is_available():
	free_before, _ = torch.cuda.mem_get_info()
	print(f"Free before: {free_before / 1024**3:.2f} GB")
	
	t = torch.randn(1024, 1024, 1024).cuda()
	
	free_after, _ = torch.cuda.mem_get_info()
	print(f"Free after: {free_after / 1024**3:.2f} GB")
	print(f"Used: {(free_before - free_after) / 1024**3:.2f} GB")

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

Free before: 11.72 GB
Free after: 7.73 GB
Used: 3.99 GB

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

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