MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / LazyNumpyArray

Class LazyNumpyArray

tools/Polygraphy/polygraphy/comparator/struct.py:28–62  ·  view source on GitHub ↗

Represents a lazily loaded NumPy array. For example, large NumPy arrays may be serialized to temporary files on the disk to save memory.

Source from the content-addressed store, hash-verified

26
27
28class LazyNumpyArray:
29 """
30 Represents a lazily loaded NumPy array.
31 For example, large NumPy arrays may be serialized to temporary files on the disk
32 to save memory.
33 """
34
35 def __init__(self, arr):
36 """
37 Args:
38 arr (np.ndarray): The NumPy array.
39 """
40 self.arr = None
41 self.tmpfile = None
42 if config.ARRAY_SWAP_THRESHOLD_MB >= 0 and arr.nbytes > (config.ARRAY_SWAP_THRESHOLD_MB << 20):
43 self.tmpfile = util.NamedTemporaryFile(suffix=".json")
44 G_LOGGER.extra_verbose(
45 f"Evicting large array ({arr.nbytes / 1024.0 ** 2:.3f} MiB) from memory and saving to {self.tmpfile.name}"
46 )
47 save_json(arr, self.tmpfile.name)
48 else:
49 self.arr = arr
50
51 def numpy(self):
52 """
53 Get the NumPy array, deserializing from the disk if it was stored earlier.
54
55 Returns:
56 np.ndarray: The NumPy array
57 """
58 if self.arr is not None:
59 return self.arr
60
61 assert self.tmpfile is not None, "Path and NumPy array cannot both be None!"
62 return load_json(self.tmpfile.name)
63
64
65@Encoder.register(LazyNumpyArray)

Callers 4

test_unswapped_arrayMethod · 0.90
test_swapped_arrayMethod · 0.90
decodeFunction · 0.85
_to_lazyMethod · 0.85

Calls

no outgoing calls

Tested by 2

test_unswapped_arrayMethod · 0.72
test_swapped_arrayMethod · 0.72