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.