Spaces:
Running on Zero
Running on Zero
File size: 3,263 Bytes
3d736fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import threading
from typing import Dict, List, Optional
import torch
class PinMemoryCache:
force_dtype: Optional[torch.dtype] = None
min_cache_numel: int = 0
pre_alloc_numels: List[int] = []
def __init__(self):
self.cache: Dict[int, torch.Tensor] = {}
self.output_to_cache: Dict[int, int] = {}
self.cache_to_output: Dict[int, int] = {}
self.lock = threading.Lock()
self.total_cnt = 0
self.hit_cnt = 0
if len(self.pre_alloc_numels) > 0 and self.force_dtype is not None:
for n in self.pre_alloc_numels:
cache_tensor = torch.empty(n, dtype=self.force_dtype, device="cpu", pin_memory=True)
with self.lock:
self.cache[id(cache_tensor)] = cache_tensor
def get(self, tensor: torch.Tensor) -> torch.Tensor:
"""Receive a cpu tensor and return the corresponding pinned tensor. Note that this only manage memory allocation, doesn't copy content.
Args:
tensor (torch.Tensor): The tensor to be pinned.
Returns:
torch.Tensor: The pinned tensor.
"""
self.total_cnt += 1
with self.lock:
# find free cache
for cache_id, cache_tensor in self.cache.items():
if cache_id not in self.cache_to_output and cache_tensor.numel() >= tensor.numel():
target_cache_tensor = cache_tensor[: tensor.numel()].view(tensor.shape)
out_id = id(target_cache_tensor)
self.output_to_cache[out_id] = cache_id
self.cache_to_output[cache_id] = out_id
self.hit_cnt += 1
return target_cache_tensor
# no free cache, create a new one
dtype = self.force_dtype if self.force_dtype is not None else tensor.dtype
cache_numel = max(tensor.numel(), self.min_cache_numel)
cache_tensor = torch.empty(cache_numel, dtype=dtype, device="cpu", pin_memory=True)
target_cache_tensor = cache_tensor[: tensor.numel()].view(tensor.shape)
out_id = id(target_cache_tensor)
with self.lock:
self.cache[id(cache_tensor)] = cache_tensor
self.output_to_cache[out_id] = id(cache_tensor)
self.cache_to_output[id(cache_tensor)] = out_id
return target_cache_tensor
def remove(self, output_tensor: torch.Tensor) -> None:
"""Release corresponding cache tensor.
Args:
output_tensor (torch.Tensor): The tensor to be released.
"""
out_id = id(output_tensor)
with self.lock:
if out_id not in self.output_to_cache:
raise ValueError("Tensor not found in cache.")
cache_id = self.output_to_cache.pop(out_id)
del self.cache_to_output[cache_id]
def __str__(self):
with self.lock:
num_cached = len(self.cache)
num_used = len(self.output_to_cache)
total_cache_size = sum([v.numel() * v.element_size() for v in self.cache.values()])
return f"PinMemoryCache(num_cached={num_cached}, num_used={num_used}, total_cache_size={total_cache_size / 1024**3:.2f} GB, hit rate={self.hit_cnt / self.total_cnt:.2f})"
|