Memory

Memory modules for benchmarking

---

source

compute_memory_multi

def compute_memory_multi(
    model:torch.nn.Module, # model to benchmark
    sample:torch.Tensor, # input tensor (with batch dimension)
    devices:Sequence[str | torch.device] | None=None, # devices to benchmark (default: cpu + cuda)
    warmup:int=10, # warmup iterations
    steps:int=100, # measurement iterations
)->dict[str, MemoryMetrics]:

Measure memory on multiple devices.

---

source

compute_memory

def compute_memory(
    model:torch.nn.Module, # model to benchmark
    sample:torch.Tensor, # input tensor (with batch dimension)
    device:str | torch.device='cpu', # device to run on
    warmup:int=10, # warmup iterations
    steps:int=100, # measurement iterations
)->MemoryMetrics:

Measure memory usage on specified device.

---

source

MemoryMetrics

def MemoryMetrics(
    avg_mib:float, peak_mib:float, reserved_mib:float, device:str
)->None:

Memory usage metrics for a single device.