Memory
Memory modules for benchmarking
compute_memory_multi
compute_memory_multi (model:torch.nn.modules.module.Module, dummy_input:torch.Tensor, devices:Optional[Iterable[str]]=None, warmup_rounds:int=10, test_rounds:int=100)
Return memory metrics for each device in devices* (default cpu + cuda).*
compute_memory
compute_memory (model:torch.nn.modules.module.Module, dummy_input:torch.Tensor, warmup_rounds:int=10, test_rounds:int=100)
Alias for GPU if available, else CPU.
MemoryMetrics
MemoryMetrics (avg_mib:float, peak_mib:float, reserved_mib:float, device:str)
Average & peak resident memory (MiB) for one device.