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.