Memory

Memory modules for benchmarking

source

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).*


source

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.


source

MemoryMetrics

 MemoryMetrics (avg_mib:float, peak_mib:float, reserved_mib:float,
                device:str)

Average & peak resident memory (MiB) for one device.