Функция empty_cache

Функция empty_cache из модуля torch.cuda освобождает всю неиспользуемую память, которая удерживается кэшем распределителя CUDA. Эта функция помогает уменьшить занимаемую видеопамять, но не гарантирует, что память будет доступна другим приложениям.

Она не освобождает память, занятую живыми тензорами, и не очищает сам кэш распределителя. Функция полезна для снижения пикового использования памяти и может быть вызвана перед запуском операций, требующих большого количества свободной видеопамяти.

Синтаксис

torch.cuda.empty_cache()

Пример использования

Давайте создадим большой тензор на GPU, удалим его и затем очистим кэш:

import torch

# Создаём большой тензор на GPU
t = torch.randn(10000, 10000, device="cuda")
print(f"Размер тензора: {t.element_size() * t.numel() / 1024**2:.2f} MB")

# Удаляем тензор
del t

# Очищаем кэш
torch.cuda.empty_cache()

# Проверяем занятую память
allocated = torch.cuda.memory_allocated()
print(f"Занято после очистки: {allocated / 1024**2:.2f} MB")

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

"Размер тензора: 381.47 MB"
"Занято после очистки: 0.00 MB"

Пример с несколькими тензорами

Рассмотрим ситуацию, когда часть тензоров всё ещё используется, а часть уже не нужна:

import torch

t1 = torch.randn(5000, 5000, device="cuda")
t2 = torch.randn(5000, 5000, device="cuda")
t3 = torch.randn(5000, 5000, device="cuda")

# Удаляем два тензора
del t1
del t2

torch.cuda.empty_cache()

# t3 всё ещё занимает память
allocated = torch.cuda.memory_allocated()
print(f"Занято после очистки: {allocated / 1024**2:.2f} MB")

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

"Занято после очистки: 95.37 MB"

Пример с изменением размера кэша

Очистка кэша может быть полезна перед выполнением операций с непредсказуемым использованием памяти:

import torch

torch.manual_seed(0)

# Создаём большой тензор и удаляем его
t = torch.randn(8000, 8000, device="cuda")
del t

# Перед очисткой
before = torch.cuda.memory_allocated()
print(f"До очистки: {before / 1024**2:.2f} MB")

torch.cuda.empty_cache()

# После очистки
after = torch.cuda.memory_allocated()
print(f"После очистки: {after / 1024**2:.2f} MB")

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

"До очистки: 0.00 MB"
"После очистки: 0.00 MB"

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

функцию memory_allocated,
которая возвращает объём занятой памяти
функцию synchronize,
которая синхронизирует все потоки CUDA
функцию reset_peak_memory_stats,
которая сбрасывает пиковые значения памяти
функцию memory_stats,
которая возвращает детальную статистику памяти