mirror of
https://github.com/brycedrennan/imaginAIry
synced 2024-10-31 03:20:40 +00:00
381 lines
12 KiB
Python
381 lines
12 KiB
Python
"""Utilities for managing model caching"""
|
|
|
|
# pylama: ignore=W0212
|
|
import logging
|
|
from collections import OrderedDict
|
|
from functools import cached_property
|
|
|
|
from imaginairy.utils import get_device
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
log = logger.debug
|
|
|
|
|
|
def get_model_size(model):
|
|
from torch import nn
|
|
|
|
if not isinstance(model, nn.Module) and hasattr(model, "model"):
|
|
model = model.model
|
|
|
|
return sum(v.numel() * v.element_size() for v in model.parameters())
|
|
|
|
|
|
def move_model_device(model, device):
|
|
from torch import nn
|
|
|
|
if not isinstance(model, nn.Module) and hasattr(model, "model"):
|
|
model = model.model
|
|
|
|
return model.to(device)
|
|
|
|
|
|
class MemoryTrackingCache:
|
|
def __init__(self, *args, **kwargs):
|
|
self.memory_usage = 0
|
|
self._item_memory_usage = {}
|
|
self._cache = OrderedDict()
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def first_key(self):
|
|
if self._cache:
|
|
return next(iter(self._cache))
|
|
raise KeyError("Empty dictionary")
|
|
|
|
def last_key(self):
|
|
if self._cache:
|
|
return next(reversed(self._cache))
|
|
raise KeyError("Empty dictionary")
|
|
|
|
def set(self, key, value, memory_usage=None):
|
|
if key in self._cache:
|
|
# Subtract old item memory usage if key already exists
|
|
self.memory_usage -= self._item_memory_usage[key]
|
|
|
|
self._cache[key] = value
|
|
|
|
# Calculate and store new item memory usage
|
|
item_memory_usage = max(get_model_size(value), memory_usage)
|
|
self._item_memory_usage[key] = item_memory_usage
|
|
self.memory_usage += item_memory_usage
|
|
|
|
def pop(self, key):
|
|
# Subtract item memory usage before deletion
|
|
self.memory_usage -= self._item_memory_usage[key]
|
|
del self._item_memory_usage[key]
|
|
return self._cache.pop(key)
|
|
|
|
def move_to_end(self, key, last=True):
|
|
self._cache.move_to_end(key, last=last)
|
|
|
|
def __contains__(self, item):
|
|
return item in self._cache
|
|
|
|
def __delitem__(self, key):
|
|
self.pop(key)
|
|
|
|
def __getitem__(self, item):
|
|
return self._cache[item]
|
|
|
|
def get(self, item):
|
|
return self._cache.get(item)
|
|
|
|
def __len__(self):
|
|
return len(self._cache)
|
|
|
|
def __bool__(self):
|
|
return bool(self._cache)
|
|
|
|
def keys(self):
|
|
return self._cache.keys()
|
|
|
|
|
|
def get_mem_free_total(device):
|
|
import psutil
|
|
import torch
|
|
|
|
if device.type == "cuda":
|
|
if not torch.cuda.is_initialized():
|
|
torch.cuda.init()
|
|
stats = torch.cuda.memory_stats(device)
|
|
mem_active = stats["active_bytes.all.current"]
|
|
mem_reserved = stats["reserved_bytes.all.current"]
|
|
mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
|
|
mem_free_torch = mem_reserved - mem_active
|
|
mem_free_total = mem_free_cuda + mem_free_torch
|
|
mem_free_total *= 0.9
|
|
else:
|
|
# if we don't add a buffer, larger images come out as noise
|
|
mem_free_total = psutil.virtual_memory().available * 0.6
|
|
|
|
return mem_free_total
|
|
|
|
|
|
class GPUModelCache:
|
|
def __init__(self, max_cpu_memory_gb="80%", max_gpu_memory_gb="95%", device=None):
|
|
self._device = device
|
|
if device in ("cpu", "mps"):
|
|
# the "gpu" cache will be the only thing we use since there aren't two different memory stores in this case
|
|
max_cpu_memory_gb = 0
|
|
|
|
self._max_cpu_memory_gb = max_cpu_memory_gb
|
|
self._max_gpu_memory_gb = max_gpu_memory_gb
|
|
self.gpu_cache = MemoryTrackingCache()
|
|
self.cpu_cache = MemoryTrackingCache()
|
|
|
|
def stats_msg(self):
|
|
import psutil
|
|
|
|
msg = (
|
|
f" GPU cache: {len(self.gpu_cache)} items; {self.gpu_cache.memory_usage / (1024 ** 2):.1f} MB; Max: {self.max_gpu_memory / (1024 ** 2):.1f} MB;\n"
|
|
f" CPU cache: {len(self.cpu_cache)} items; {self.cpu_cache.memory_usage / (1024 ** 2):.1f} MB; Max: {self.max_cpu_memory / (1024 ** 2):.1f} MB;\n"
|
|
f" mem_free_total: {get_mem_free_total(self.device) / (1024 ** 2):.1f} MB; Ram Free: {psutil.virtual_memory().available / (1024 ** 2):.1f} MB;"
|
|
)
|
|
return msg
|
|
|
|
@cached_property
|
|
def device(self):
|
|
import torch
|
|
|
|
if self._device is None:
|
|
self._device = get_device()
|
|
|
|
if self._device in ("cpu", "mps", "mps:0"):
|
|
# the "gpu" cache will be the only thing we use since there aren't two different memory stores in this case
|
|
self._max_cpu_memory_gb = 0
|
|
|
|
return torch.device(self._device)
|
|
|
|
def make_gpu_space(self, bytes_to_free):
|
|
import gc
|
|
|
|
import torch.cuda
|
|
|
|
log(self.stats_msg())
|
|
log(f"Ensuring {bytes_to_free / (1024 ** 2):.1f} MB of GPU space.")
|
|
|
|
while self.gpu_cache and (
|
|
self.gpu_cache.memory_usage + bytes_to_free > self.max_gpu_memory
|
|
or self.gpu_cache.memory_usage + bytes_to_free
|
|
> get_mem_free_total(self.device)
|
|
):
|
|
oldest_gpu_key = self.gpu_cache.first_key()
|
|
self._move_to_cpu(oldest_gpu_key)
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
gc.collect()
|
|
|
|
if (
|
|
self.gpu_cache.memory_usage + bytes_to_free > self.max_gpu_memory
|
|
or self.gpu_cache.memory_usage + bytes_to_free
|
|
> get_mem_free_total(self.device)
|
|
):
|
|
msg = f"Unable to make {bytes_to_free / (1024 ** 2):.1f} MB space on {self.device}. \n{self.stats_msg()}"
|
|
raise RuntimeError(msg)
|
|
|
|
def make_cpu_space(self, bytes_to_free):
|
|
import gc
|
|
|
|
import psutil
|
|
|
|
log(self.stats_msg())
|
|
log(f"Ensuring {bytes_to_free / (1024 ** 2):.1f} MB of RAM space.")
|
|
while self.cpu_cache and (
|
|
self.cpu_cache.memory_usage + bytes_to_free > self.max_gpu_memory
|
|
or self.cpu_cache.memory_usage + bytes_to_free
|
|
> psutil.virtual_memory().available * 0.8
|
|
):
|
|
oldest_cpu_key = self.cpu_cache.first_key()
|
|
log(f"dropping {oldest_cpu_key} from memory")
|
|
self.cpu_cache.pop(oldest_cpu_key)
|
|
log(self.stats_msg())
|
|
|
|
gc.collect()
|
|
|
|
@cached_property
|
|
def max_cpu_memory(self):
|
|
_ = self.device
|
|
if isinstance(self._max_cpu_memory_gb, str):
|
|
if self._max_cpu_memory_gb.endswith("%"):
|
|
import psutil
|
|
|
|
total_ram_gb = round(psutil.virtual_memory().total / (1024**3), 2)
|
|
pct_to_use = float(self._max_cpu_memory_gb[:-1]) / 100.0
|
|
return total_ram_gb * pct_to_use * (1024**3)
|
|
msg = f"Invalid value for max_cpu_memory_gb: {self._max_cpu_memory_gb}"
|
|
raise ValueError(msg)
|
|
return self._max_cpu_memory_gb * (1024**3)
|
|
|
|
@cached_property
|
|
def max_gpu_memory(self):
|
|
_ = self.device
|
|
if isinstance(self._max_gpu_memory_gb, str):
|
|
if self._max_gpu_memory_gb.endswith("%"):
|
|
import torch
|
|
|
|
if self.device.type == "cuda":
|
|
device_props = torch.cuda.get_device_properties(0)
|
|
total_ram_gb = device_props.total_memory / (1024**3)
|
|
else:
|
|
import psutil
|
|
|
|
total_ram_gb = round(psutil.virtual_memory().total / (1024**3), 2)
|
|
pct_to_use = float(self._max_gpu_memory_gb[:-1]) / 100.0
|
|
return total_ram_gb * pct_to_use * (1024**3)
|
|
msg = f"Invalid value for max_gpu_memory_gb: {self._max_gpu_memory_gb}"
|
|
raise ValueError(msg)
|
|
return self._max_gpu_memory_gb * (1024**3)
|
|
|
|
def _move_to_gpu(self, key, model):
|
|
model_size = get_model_size(model)
|
|
if self.gpu_cache.memory_usage + model_size > self.max_gpu_memory:
|
|
if len(self.gpu_cache) == 0:
|
|
msg = f"GPU cache maximum ({self.max_gpu_memory / (1024 ** 2)} MB) is smaller than the item being cached ({model_size / 1024 ** 2} MB)."
|
|
raise RuntimeError(msg)
|
|
self.make_gpu_space(model_size)
|
|
|
|
try:
|
|
model_size = max(self.cpu_cache._item_memory_usage[key], model_size)
|
|
self.cpu_cache.pop(key)
|
|
log(f"dropping {key} from cpu cache")
|
|
except KeyError:
|
|
pass
|
|
log(f"moving {key} to gpu")
|
|
move_model_device(model, self.device)
|
|
|
|
self.gpu_cache.set(key, value=model, memory_usage=model_size)
|
|
|
|
def _move_to_cpu(self, key):
|
|
import gc
|
|
|
|
import psutil
|
|
import torch
|
|
|
|
memory_usage = self.gpu_cache._item_memory_usage[key]
|
|
model = self.gpu_cache.pop(key)
|
|
|
|
model_size = max(get_model_size(model), memory_usage)
|
|
self.make_cpu_space(model_size)
|
|
|
|
if (
|
|
self.cpu_cache.memory_usage + model_size < self.max_cpu_memory
|
|
and self.cpu_cache.memory_usage + model_size
|
|
< psutil.virtual_memory().available * 0.8
|
|
):
|
|
log(f"moving {key} to cpu")
|
|
move_model_device(model, torch.device("cpu"))
|
|
log(self.stats_msg())
|
|
|
|
self.cpu_cache.set(key, model, memory_usage=model_size)
|
|
else:
|
|
log(f"dropping {key} from memory")
|
|
del model
|
|
gc.collect()
|
|
log(self.stats_msg())
|
|
|
|
def get(self, key):
|
|
import torch
|
|
|
|
if key not in self:
|
|
msg = f"The key {key} does not exist in the cache"
|
|
raise KeyError(msg)
|
|
|
|
if key in self.cpu_cache and self.device != torch.device("cpu"):
|
|
self.cpu_cache.move_to_end(key)
|
|
self._move_to_gpu(key, self.cpu_cache[key])
|
|
|
|
if key in self.gpu_cache:
|
|
self.gpu_cache.move_to_end(key)
|
|
|
|
model = self.gpu_cache.get(key)
|
|
|
|
return model
|
|
|
|
def __getitem__(self, key):
|
|
return self.get(key)
|
|
|
|
def set(self, key, model, memory_usage=0):
|
|
from torch import nn
|
|
|
|
if (
|
|
hasattr(model, "model") and isinstance(model.model, nn.Module)
|
|
) or isinstance(model, nn.Module):
|
|
pass
|
|
else:
|
|
raise ValueError("Only nn.Module objects can be cached")
|
|
|
|
model_size = max(get_model_size(model), memory_usage)
|
|
self.make_gpu_space(model_size)
|
|
self._move_to_gpu(key, model)
|
|
|
|
def __contains__(self, key):
|
|
return key in self.gpu_cache or key in self.cpu_cache
|
|
|
|
def keys(self):
|
|
return list(self.cpu_cache.keys()) + list(self.gpu_cache.keys())
|
|
|
|
def stats(self):
|
|
return {
|
|
"cpu_cache_count": len(self.cpu_cache),
|
|
"cpu_cache_memory_usage": self.cpu_cache.memory_usage,
|
|
"cpu_cache_max_memory": self.max_cpu_memory,
|
|
"gpu_cache_count": len(self.gpu_cache),
|
|
"gpu_cache_memory_usage": self.gpu_cache.memory_usage,
|
|
"gpu_cache_max_memory": self.max_gpu_memory,
|
|
}
|
|
|
|
|
|
class MemoryManagedModelWrapper:
|
|
_mmmw_cache = GPUModelCache()
|
|
|
|
def __init__(self, fn, namespace, estimated_ram_size_mb, *args, **kwargs):
|
|
self._mmmw_fn = fn
|
|
self._mmmw_args = args
|
|
self._mmmw_kwargs = kwargs
|
|
self._mmmw_namespace = namespace
|
|
self._mmmw_estimated_ram_size_mb = estimated_ram_size_mb
|
|
self._mmmw_cache_key = (namespace, *args, *tuple(kwargs.items()))
|
|
|
|
def _mmmw_load_model(self):
|
|
if self._mmmw_cache_key not in self.__class__._mmmw_cache:
|
|
log(f"Loading model: {self._mmmw_cache_key}")
|
|
self.__class__._mmmw_cache.make_gpu_space(
|
|
self._mmmw_estimated_ram_size_mb * 1024**2
|
|
)
|
|
free_before = get_mem_free_total(self.__class__._mmmw_cache.device)
|
|
model = self._mmmw_fn(*self._mmmw_args, **self._mmmw_kwargs)
|
|
move_model_device(model, self.__class__._mmmw_cache.device)
|
|
free_after = get_mem_free_total(self.__class__._mmmw_cache.device)
|
|
log(f"Model loaded: {self._mmmw_cache_key} Used {free_after - free_before}")
|
|
self.__class__._mmmw_cache.set(
|
|
self._mmmw_cache_key,
|
|
model,
|
|
memory_usage=self._mmmw_estimated_ram_size_mb * 1024**2,
|
|
)
|
|
|
|
model = self.__class__._mmmw_cache[self._mmmw_cache_key]
|
|
return model
|
|
|
|
def __getattr__(self, name):
|
|
model = self._mmmw_load_model()
|
|
return getattr(model, name)
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
model = self._mmmw_load_model()
|
|
return model(*args, **kwargs)
|
|
|
|
|
|
def memory_managed_model(namespace, memory_usage_mb=0):
|
|
def decorator(fn):
|
|
def wrapper(*args, **kwargs):
|
|
return MemoryManagedModelWrapper(
|
|
fn, namespace, memory_usage_mb, *args, **kwargs
|
|
)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|