Compression Utilities¶

auto_set_denpendency_group_ids¶

class nni.compression.utils.auto_set_denpendency_group_ids(model, config_list, dummy_input)[source]¶

Auto find the output dependency between all ‘Conv2d’, ‘Linear’, ‘ConvTranspose2d’, ‘Embedding’ modules, then set the dependency_group_id in config list.

Note that a new dependency group id will be set as a shortcut in one config, it will replace the old configured one in that config.

Parameters:

model (torch.nn.Module) – The origin model.
config_list (List[Dict[str, Any]]) – The compression config list.
dummy_input (torch.Tensor | List[torch.Tensor] | Dict[str, torch.Tensor]) – The dummy input to the model forward function for tracing the model.