PyTorch 2.1.0 and Lightning 2.1.0 Support: AssertionError on `assert self._strategy is not None` about sparseml HOT 13 CLOSED

    def on_fit_start(self, trainer: Trainer,
                     pl_module: LightningModule) -> None:
        if len(trainer.strategy.optimizers) > 1:
            raise MisconfigurationException(
                "SparseML only supports training with one optimizer.")
        optimizer = trainer.strategy.optimizers[0]

        optimizer = self.manager.modify(
            pl_module,
            optimizer,
            steps_per_epoch=trainer.estimated_stepping_batches,
            epoch=0)

        trainer.strategy._optimizers = [optimizer]
        trainer.strategy._lightning_optimizers[0]._optimizer = optimizer

from typing import Any, Optional

import torch
from lightning.pytorch import Callback, LightningModule, Trainer
from lightning.pytorch.utilities.exceptions import MisconfigurationException
from sparseml.pytorch.optim import ScheduledModifierManager
from sparseml.pytorch.utils import ModuleExporter


class SparseMLCallback(Callback):
    """Enables SparseML aware training. Requires a recipe to run during training.

    Args:
        recipe_path: Path to a SparseML compatible yaml recipe.
            More information at https://docs.neuralmagic.com/sparseml/source/recipes.html

    """

    def __init__(self,
                 recipe_path: str,
                 steps_per_epoch: Optional[int] = None) -> None:
        self.manager = ScheduledModifierManager.from_yaml(recipe_path)
        self.steps_per_epoch = steps_per_epoch

    def on_fit_start(self, trainer: Trainer,
                     pl_module: LightningModule) -> None:
        optimizer = trainer.strategy._lightning_optimizers[0]._optimizer

        if len(trainer.optimizers) > 1:
            raise MisconfigurationException(
                "SparseML only supports training with one optimizer.")

        if self.steps_per_epoch is None:
            self.steps_per_epoch = trainer.estimated_stepping_batches

        optimizer = self.manager.modify(pl_module,
                                        optimizer,
                                        steps_per_epoch=self.steps_per_epoch,
                                        epoch=0)

        trainer.strategy._lightning_optimizers[0]._optimizer = optimizer

    def on_fit_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
        self.manager.finalize(pl_module)

    @staticmethod
    # TODO: check for TrainingMode.EVAL
    def export_to_sparse_onnx(model: LightningModule,
                              output_dir: str,
                              sample_batch: Optional[torch.Tensor] = None,
                              name: str = "model.onnx",
                              opset: int = 14,
                              disable_bn_fusing: bool = True,
                              convert_qat: bool = True,
                              **export_kwargs: Any) -> None:
        """Exports the model to ONNX format."""
        exporter = ModuleExporter(model, output_dir=output_dir)
        sample_batch = sample_batch if sample_batch is not None else model.example_input_array  # type: ignore[assignment] # noqa: E501
        if sample_batch is None:
            raise MisconfigurationException(
                "To export the model, a sample batch must be passed via "
                "``SparseMLCallback.export_to_sparse_onnx(model, output_dir, sample_batch=sample_batch)`` "
                "or an ``example_input_array`` property within the LightningModule"
            )
        exporter.export_onnx(
            sample_batch=sample_batch,
            name=name,
            opset=opset,
            disable_bn_fusing=disable_bn_fusing,
            convert_qat=convert_qat,
            **export_kwargs,
        )

# Export to ONNX
clf.eval()
clf.training = False
sparseml.export_to_sparse_onnx(
    clf,
    output_dir="./sparseml_models/",
    name="sparse_model.onnx",
    sample_batch=torch.randn(1, 1, 28, 28),
    # opset_version=17,
    opset=14,
    disable_bn_fusing=False,
    convert_qat=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {
            0: "batch_size"
        },
        "output": {
            0: "batch_size"
        },
    },
)