Benchmark

Benchmark module

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’) subset of {“size”,“speed”,“compute”,“memory”,“energy”} to compute
speed_devices Optional None per-family device override; Nonemeans “CPU and CUDA if available”
memory_devices Optional None
energy_devices Optional None
kwargs VAR_KEYWORD
Returns Dict