Benchmark

Unified benchmarking API for PyTorch models

source

benchmark

 benchmark (model:torch.nn.modules.module.Module, sample:torch.Tensor,
            metrics:Sequence[str]=('size', 'speed', 'compute', 'memory',
            'energy'),
            speed_devices:Optional[Sequence[str|torch.device]]=None,
            memory_devices:Optional[Sequence[str|torch.device]]=None,
            energy_devices:Optional[Sequence[str|torch.device]]=None,
            **kwargs)
Type Default Details
model Module the model to profile (can stay on CPU)
sample Tensor dummy input of the right shape
metrics Sequence (‘size’, ‘speed’, ‘compute’, ‘memory’, ‘energy’)
speed_devices Optional None
memory_devices Optional None
energy_devices Optional None
kwargs
Returns Dict

See Also