forked from arikpoz/deep-visualization-toolbox
-
Notifications
You must be signed in to change notification settings - Fork 0
/
numpy_cache.py
50 lines (42 loc) · 1.69 KB
/
numpy_cache.py
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
import numpy as np
from collections import OrderedDict
from threading import RLock
class FIFOLimitedArrayCache(object):
'''Threadsafe cache that stores numpy arrays (or any other object that
defines obj.nbytes) and limits total memory. Items are ejected, if
necessary, in the same order in which they were added.
'''
def __init__(self, max_bytes = 1e7):
self._store = OrderedDict()
self._store_bytes = 0
self._max_bytes = max_bytes
self._lock = RLock()
def get(self, key, default = None):
with self._lock:
if key in self._store:
return self._store[key]
else:
return default
def set(self, key, val):
with self._lock:
if key in self._store:
self._store_bytes -= self._store[key].nbytes
self._store[key] = val
self._store_bytes += self._store[key].nbytes
self._trim()
def _trim(self):
while len(self._store) > 0 and self._store_bytes > self._max_bytes:
key,val = self._store.popitem(last = False)
self._store_bytes -= val.nbytes
def delete(self, key, raise_if_missing = False):
with self._lock:
if key in self._store:
self._store_bytes -= val.nbytes
del self._store[key]
elif raise_if_missing:
raise Exception('key %s not found in cache' % repr(key))
def get_size(self):
return self._store_bytes
def __str__(self):
with self._lock:
return 'FIFOLimitedArrayCache<%d items, bytes used/max %g/%g >' % (len(self._store), self._store_bytes, self._max_bytes)