llmcompressor.modifiers.obcq.sgpt_base

def on_initialize(self, state: "State", **kwargs) -> bool:
    """
    Initialize and run the OBCQ algorithm on the current state

    :param state: session state storing input model and calibration data
    """
    model: torch.nn.Module = state.model
    dataloader: torch.utils.data.DataLoader = state.data.calib

    # infer module and sequential targets
    self.sequential_targets = self._infer_sequential_targets(model)
    layers = get_layers(self.sequential_targets, model)
    self._target_layers = get_layers(
        self.targets, model
    )  # layers containing targets

    # infer layer sparsities
    if self.sparsity_profile == "owl":
        logger.info(
            "Using OWL to infer target layer-wise sparsities from "
            f"{len(dataloader) if dataloader else 0} calibration samples..."
        )
        self.sparsity = self._infer_owl_layer_sparsity(model, layers, dataloader)

    # get layers and validate sparsity
    if isinstance(self.sparsity, (list, dict)) and len(self._target_layers) != len(
        self.sparsity
    ):
        raise ValueError(
            f"{self.__repr_name__} was initialized with {len(self.sparsity)} "
            f"sparsities values, but model has {len(layers)} target layers"
        )

    return True