dmx.compressor.utils.benchmark.collect_layer_activations

dmx.compressor.utils.benchmark.collect_layer_activations(model_maker: Callable[[], Tuple[Module, Callable, Callable, device]], mode: EVALUATION_MODE)

Collect the output activations for each named DMX module in the model for the given mode

Parameters:

• model_maker (Callable[[], Tuple[torch.nn.Module, Callable, Callable, torch.device]]) – A callable that returns the model to be measured, together with some callables to run a sample input through the model or to evaluate the model’s accuracy

• mode (EVALUATION_MODE) – The mode for which we want to collect the layer activations