Compression Utilities¶
auto_set_denpendency_group_ids¶
- class nni.contrib.compression.utils.auto_set_denpendency_group_ids(model, config_list, dummy_input)[源代码]¶
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.
- 参数:
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.