distillation

shrinkai.distillation

Knowledge distillation: training a small student model to mimic a larger teacher.

Distiller is the high-level facade most users start from: give it a teacher, a student, and a loss (shrinkai.distillation.losses), and it handles the training loop (fit), evaluation, benchmarking, checkpointing, and deployment export. DistillationEngine is the lower-level training loop it delegates to (device management, mixed precision, gradient clipping, teacher freezing, ...), usable directly for custom orchestration. EarlyStopping/ModelCheckpoint are ready-to-use fit(callbacks=[...]) callbacks.

Modules:

Name	Description
callbacks	Ready-to-use training callbacks for DistillationEngine.fit / Distiller.fit.
distiller
engine
losses	Knowledge Distillation loss functions.

Classes:

Name	Description
DistillationEngine	Generic training engine for knowledge distillation across hardware targets.
Distiller	Unified, high-level facade for end-to-end knowledge distillation and benchmarking.
EarlyStopping	Stops training when a monitored metric has stopped improving.
ModelCheckpoint	Saves a model's weights to disk during training.

Classes

DistillationEngine

Generic training engine for knowledge distillation across hardware targets.

Handles training/validation loops, accelerator management (MPS, CUDA, CPU), teacher state freezing, and metrics tracking.

Methods:

Name	Description
__init__	Initializes the DistillationEngine.
evaluate	Evaluates student performance on validation/test data.
fit	Executes the full distillation training loop.
train_epoch	Runs a single training epoch over the provided dataloader.

Source code in src/shrinkai/distillation/engine.py

class DistillationEngine:
    """Generic training engine for knowledge distillation across hardware targets.

    Handles training/validation loops, accelerator management (MPS, CUDA, CPU),
    teacher state freezing, and metrics tracking.
    """

    def __init__(
        self,
        student: nn.Module,
        teacher: nn.Module,
        criterion: BaseDistillationLoss,
        optimizer: torch.optim.Optimizer,
        device: torch.device | str = "auto",
        scheduler: torch.optim.lr_scheduler._LRScheduler | None = None,
        use_amp: bool = False,
        grad_clip_norm: float | None = None,
    ) -> None:
        """Initializes the DistillationEngine.

        Args:
            student: Student neural network module to train.
            teacher: Pre-trained Teacher neural network module providing soft targets.
            criterion: Loss function adhering to `BaseDistillationLoss`.
            optimizer: PyTorch optimizer targeting student parameters.
            device: Computing device ('auto', 'mps', 'cuda', 'cpu' or torch.device).
            scheduler: Optional learning rate scheduler updated per epoch.
            use_amp: If True, runs the forward passes and loss computation under
                mixed precision (`torch.autocast`). Uses fp16 with gradient scaling
                on CUDA, and bf16 (no scaling needed) on CPU/MPS. Defaults to False.
            grad_clip_norm: If set, clips the student's gradient global L2 norm to
                this value before each optimizer step. Defaults to None (no clipping).
        """
        self.device = resolve_device(device)
        self.student = student.to(self.device)
        self.teacher = teacher.to(self.device)
        self.criterion = criterion.to(self.device)
        self.optimizer = optimizer
        self.scheduler = scheduler

        self.use_amp = use_amp
        self.grad_clip_norm = grad_clip_norm
        self._amp_dtype = torch.float16 if self.device.type == "cuda" else torch.bfloat16
        self._use_scaler = self.use_amp and self.device.type == "cuda"
        self.scaler = torch.amp.GradScaler(device="cuda", enabled=self._use_scaler)

        self.teacher.eval()
        for param in self.teacher.parameters():
            param.requires_grad = False

    def _unpack_batch(self, batch: Any) -> tuple[Any, torch.Tensor]:
        """Unpacks tuple/list or dict batches and moves them to the device.

        Provides compatibility with both standard PyTorch Datasets (tuples)
        and Hugging Face Datasets (dictionaries).
        """
        if isinstance(batch, dict):
            # Hugging Face style batch
            batch = {
                k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()
            }
            if "labels" not in batch:
                raise ValueError("Batch dictionary must contain a 'labels' key.")
            labels = batch.pop("labels")
            return batch, labels

        elif isinstance(batch, list | tuple):
            # Standard PyTorch style batch
            inputs, labels = batch[0].to(self.device), batch[1].to(self.device)
            return inputs, labels

        else:
            raise TypeError(f"Unsupported batch type: {type(batch)}")

    def _forward_model(self, model: nn.Module, inputs: Any) -> Any:
        """Executes forward pass adapting to the input type (kwargs or args)."""
        if isinstance(inputs, dict):
            return model(**inputs)
        return model(inputs)

    def train_epoch(
        self,
        dataloader: DataLoader,
        epoch_idx: int,
        total_epochs: int,
    ) -> dict[str, float]:
        """Runs a single training epoch over the provided dataloader.

        Args:
            dataloader: Training dataloader yielding (inputs, labels) batches.
            epoch_idx: Current 1-based epoch index.
            total_epochs: Total number of planned epochs.

        Returns:
            dict[str, float]: Aggregated training metrics (loss, accuracy).
        """
        self.student.train()
        total_loss = 0.0
        correct = 0
        total_samples = 0

        desc = f"Epoch [{epoch_idx}/{total_epochs}] Training"
        pbar = tqdm(dataloader, desc=desc, leave=False)

        for batch in pbar:
            inputs, labels = self._unpack_batch(batch)

            with (
                torch.no_grad(),
                torch.autocast(
                    device_type=self.device.type, dtype=self._amp_dtype, enabled=self.use_amp
                ),
            ):
                teacher_outputs = self._forward_model(self.teacher, inputs)

            self.optimizer.zero_grad()
            with torch.autocast(
                device_type=self.device.type, dtype=self._amp_dtype, enabled=self.use_amp
            ):
                student_outputs = self._forward_model(self.student, inputs)
                loss = self.criterion(
                    student_outputs=student_outputs,
                    teacher_outputs=teacher_outputs,
                    labels=labels,
                )

            if not math.isfinite(loss.item()):
                raise RuntimeError(
                    f"Loss diverged to {loss.item()} during training at epoch {epoch_idx}. "
                    "Check learning rate or data scaling."
                )

            if self._use_scaler:
                self.scaler.scale(loss).backward()
                if self.grad_clip_norm is not None:
                    self.scaler.unscale_(self.optimizer)
                    nn.utils.clip_grad_norm_(self.student.parameters(), self.grad_clip_norm)
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                loss.backward()
                if self.grad_clip_norm is not None:
                    nn.utils.clip_grad_norm_(self.student.parameters(), self.grad_clip_norm)
                self.optimizer.step()

            student_logits = (
                student_outputs[0] if isinstance(student_outputs, tuple) else student_outputs
            )

            if student_logits.dim() == 3:
                # LLM Causal Shift
                preds = torch.argmax(student_logits[..., :-1, :].contiguous(), dim=-1)
                shifted_labels = labels[..., 1:].contiguous()
                correct += (preds == shifted_labels).sum().item()
                batch_samples = shifted_labels.numel()
            else:
                preds = torch.argmax(student_logits, dim=-1)
                correct += (preds == labels).sum().item()
                batch_samples = labels.size(0)

            total_loss += loss.item() * batch_samples
            total_samples += batch_samples

            current_loss = total_loss / total_samples
            current_acc = (correct / total_samples) * 100.0
            pbar.set_postfix(loss=f"{current_loss:.4f}", acc=f"{current_acc:.2f}%")

        return {
            "loss": total_loss / total_samples,
            "accuracy": (correct / total_samples) * 100.0,
        }

    def evaluate(self, dataloader: DataLoader) -> dict[str, float]:
        """Evaluates student performance on validation/test data.

        Args:
            dataloader: Validation dataloader yielding (inputs, labels) batches.

        Returns:
            dict[str, float]: Validation metrics (loss, accuracy).
        """
        self.student.eval()
        total_loss = 0.0
        correct = 0
        total_samples = 0

        with (
            torch.no_grad(),
            torch.autocast(
                device_type=self.device.type, dtype=self._amp_dtype, enabled=self.use_amp
            ),
        ):
            for batch in dataloader:
                inputs, labels = self._unpack_batch(batch)

                teacher_outputs = self._forward_model(self.teacher, inputs)
                student_outputs = self._forward_model(self.student, inputs)

                loss = self.criterion(
                    student_outputs=student_outputs,
                    teacher_outputs=teacher_outputs,
                    labels=labels,
                )
                student_logits = (
                    student_outputs[0] if isinstance(student_outputs, tuple) else student_outputs
                )

                if student_logits.dim() == 3:
                    # LLM Causal Shift
                    preds = torch.argmax(student_logits[..., :-1, :].contiguous(), dim=-1)
                    shifted_labels = labels[..., 1:].contiguous()
                    correct += (preds == shifted_labels).sum().item()
                    batch_samples = shifted_labels.numel()
                else:
                    preds = torch.argmax(student_logits, dim=-1)
                    correct += (preds == labels).sum().item()
                    batch_samples = labels.size(0)

                total_loss += loss.item() * batch_samples
                total_samples += batch_samples

        return {
            "val_loss": total_loss / total_samples,
            "val_accuracy": (correct / total_samples) * 100.0,
        }

    def fit(
        self,
        train_dataloader: DataLoader,
        val_dataloader: DataLoader | None = None,
        epochs: int = 10,
        callbacks: list[Callable[[int, dict[str, float]], None]] | None = None,
        start_epoch: int = 1,
        history: dict[str, list[float]] | None = None,
    ) -> dict[str, list[float]]:
        """Executes the full distillation training loop.

        Args:
            train_dataloader: Dataloader containing training dataset.
            val_dataloader: Optional dataloader for epoch-end validation.
            epochs: Total number of epochs to train up to (1-indexed, inclusive).
                Defaults to 10.
            callbacks: Optional list of callback functions triggered each epoch.
                A callback exposing a truthy `stop` attribute after being called
                (e.g. `EarlyStopping`) interrupts training at the end of that epoch.
            start_epoch: 1-based epoch index to resume training from. Defaults to 1
                (a fresh run). Used together with `history` when resuming from a
                checkpoint saved via `Distiller.save_checkpoint`.
            history: Existing training history to append to, as returned by a
                previous call to `fit`. Defaults to None (starts a fresh history).

        Returns:
            dict[str, list[float]]: Training history tracking loss and metrics,
            covering both the resumed epochs (if any) and the new ones.
        """
        if history is None:
            history = {
                "train_loss": [],
                "train_accuracy": [],
                "val_loss": [],
                "val_accuracy": [],
            }

        for epoch in range(start_epoch, epochs + 1):
            train_metrics = self.train_epoch(train_dataloader, epoch, epochs)
            history["train_loss"].append(train_metrics["loss"])
            history["train_accuracy"].append(train_metrics["accuracy"])

            val_metrics: dict[str, float] = {}
            if val_dataloader is not None:
                val_metrics = self.evaluate(val_dataloader)
                history["val_loss"].append(val_metrics["val_loss"])
                history["val_accuracy"].append(val_metrics["val_accuracy"])

            if self.scheduler is not None:
                self.scheduler.step()

            status = (
                f"Epoch [{epoch:02d}/{epochs:02d}] "
                f"Train Loss: {train_metrics['loss']:.4f} - "
                f"Train Acc: {train_metrics['accuracy']:.2f}%"
            )
            if val_metrics:
                status += (
                    f" | Val Loss: {val_metrics['val_loss']:.4f} - "
                    f"Val Acc: {val_metrics['val_accuracy']:.2f}%"
                )
            tqdm.write(status)

            if callbacks:
                epoch_summary = {**train_metrics, **val_metrics}
                for callback in callbacks:
                    callback(epoch, epoch_summary)

                if any(getattr(callback, "stop", False) for callback in callbacks):
                    tqdm.write(f"Training stopped early at epoch {epoch}/{epochs}.")
                    break

        return history

Methods:

__init__

__init__(
    student: Module,
    teacher: Module,
    criterion: BaseDistillationLoss,
    optimizer: Optimizer,
    device: device | str = "auto",
    scheduler: _LRScheduler | None = None,
    use_amp: bool = False,
    grad_clip_norm: float | None = None,
) -> None

Initializes the DistillationEngine.

Parameters:

Name	Type	Description	Default
student	Module	Student neural network module to train.	required
teacher	Module	Pre-trained Teacher neural network module providing soft targets.	required
criterion	BaseDistillationLoss	Loss function adhering to BaseDistillationLoss.	required
optimizer	Optimizer	PyTorch optimizer targeting student parameters.	required
device	device | str	Computing device ('auto', 'mps', 'cuda', 'cpu' or torch.device).	'auto'
scheduler	_LRScheduler | None	Optional learning rate scheduler updated per epoch.	None
use_amp	bool	If True, runs the forward passes and loss computation under mixed precision (torch.autocast). Uses fp16 with gradient scaling on CUDA, and bf16 (no scaling needed) on CPU/MPS. Defaults to False.	False
grad_clip_norm	float | None	If set, clips the student's gradient global L2 norm to this value before each optimizer step. Defaults to None (no clipping).	None

Source code in src/shrinkai/distillation/engine.py

def __init__(
    self,
    student: nn.Module,
    teacher: nn.Module,
    criterion: BaseDistillationLoss,
    optimizer: torch.optim.Optimizer,
    device: torch.device | str = "auto",
    scheduler: torch.optim.lr_scheduler._LRScheduler | None = None,
    use_amp: bool = False,
    grad_clip_norm: float | None = None,
) -> None:
    """Initializes the DistillationEngine.

    Args:
        student: Student neural network module to train.
        teacher: Pre-trained Teacher neural network module providing soft targets.
        criterion: Loss function adhering to `BaseDistillationLoss`.
        optimizer: PyTorch optimizer targeting student parameters.
        device: Computing device ('auto', 'mps', 'cuda', 'cpu' or torch.device).
        scheduler: Optional learning rate scheduler updated per epoch.
        use_amp: If True, runs the forward passes and loss computation under
            mixed precision (`torch.autocast`). Uses fp16 with gradient scaling
            on CUDA, and bf16 (no scaling needed) on CPU/MPS. Defaults to False.
        grad_clip_norm: If set, clips the student's gradient global L2 norm to
            this value before each optimizer step. Defaults to None (no clipping).
    """
    self.device = resolve_device(device)
    self.student = student.to(self.device)
    self.teacher = teacher.to(self.device)
    self.criterion = criterion.to(self.device)
    self.optimizer = optimizer
    self.scheduler = scheduler

    self.use_amp = use_amp
    self.grad_clip_norm = grad_clip_norm
    self._amp_dtype = torch.float16 if self.device.type == "cuda" else torch.bfloat16
    self._use_scaler = self.use_amp and self.device.type == "cuda"
    self.scaler = torch.amp.GradScaler(device="cuda", enabled=self._use_scaler)

    self.teacher.eval()
    for param in self.teacher.parameters():
        param.requires_grad = False

evaluate

evaluate(dataloader: DataLoader) -> dict[str, float]

Evaluates student performance on validation/test data.

Parameters:

Name	Type	Description	Default
dataloader	DataLoader	Validation dataloader yielding (inputs, labels) batches.	required

Returns:

Type	Description
dict[str, float]	dict[str, float]: Validation metrics (loss, accuracy).

Source code in src/shrinkai/distillation/engine.py

def evaluate(self, dataloader: DataLoader) -> dict[str, float]:
    """Evaluates student performance on validation/test data.

    Args:
        dataloader: Validation dataloader yielding (inputs, labels) batches.

    Returns:
        dict[str, float]: Validation metrics (loss, accuracy).
    """
    self.student.eval()
    total_loss = 0.0
    correct = 0
    total_samples = 0

    with (
        torch.no_grad(),
        torch.autocast(
            device_type=self.device.type, dtype=self._amp_dtype, enabled=self.use_amp
        ),
    ):
        for batch in dataloader:
            inputs, labels = self._unpack_batch(batch)

            teacher_outputs = self._forward_model(self.teacher, inputs)
            student_outputs = self._forward_model(self.student, inputs)

            loss = self.criterion(
                student_outputs=student_outputs,
                teacher_outputs=teacher_outputs,
                labels=labels,
            )
            student_logits = (
                student_outputs[0] if isinstance(student_outputs, tuple) else student_outputs
            )

            if student_logits.dim() == 3:
                # LLM Causal Shift
                preds = torch.argmax(student_logits[..., :-1, :].contiguous(), dim=-1)
                shifted_labels = labels[..., 1:].contiguous()
                correct += (preds == shifted_labels).sum().item()
                batch_samples = shifted_labels.numel()
            else:
                preds = torch.argmax(student_logits, dim=-1)
                correct += (preds == labels).sum().item()
                batch_samples = labels.size(0)

            total_loss += loss.item() * batch_samples
            total_samples += batch_samples

    return {
        "val_loss": total_loss / total_samples,
        "val_accuracy": (correct / total_samples) * 100.0,
    }

fit

fit(
    train_dataloader: DataLoader,
    val_dataloader: DataLoader | None = None,
    epochs: int = 10,
    callbacks: list[Callable[[int, dict[str, float]], None]]
    | None = None,
    start_epoch: int = 1,
    history: dict[str, list[float]] | None = None,
) -> dict[str, list[float]]

Executes the full distillation training loop.

Parameters:

Name	Type	Description	Default
train_dataloader	DataLoader	Dataloader containing training dataset.	required
val_dataloader	DataLoader | None	Optional dataloader for epoch-end validation.	None
epochs	int	Total number of epochs to train up to (1-indexed, inclusive). Defaults to 10.	10
callbacks	list[Callable[[int, dict[str, float]], None]] | None	Optional list of callback functions triggered each epoch. A callback exposing a truthy stop attribute after being called (e.g. EarlyStopping) interrupts training at the end of that epoch.	None
start_epoch	int	1-based epoch index to resume training from. Defaults to 1 (a fresh run). Used together with history when resuming from a checkpoint saved via Distiller.save_checkpoint.	1
history	dict[str, list[float]] | None	Existing training history to append to, as returned by a previous call to fit. Defaults to None (starts a fresh history).	None

Returns:

Type	Description
dict[str, list[float]]	dict[str, list[float]]: Training history tracking loss and metrics,
dict[str, list[float]]	covering both the resumed epochs (if any) and the new ones.

Source code in src/shrinkai/distillation/engine.py

def fit(
    self,
    train_dataloader: DataLoader,
    val_dataloader: DataLoader | None = None,
    epochs: int = 10,
    callbacks: list[Callable[[int, dict[str, float]], None]] | None = None,
    start_epoch: int = 1,
    history: dict[str, list[float]] | None = None,
) -> dict[str, list[float]]:
    """Executes the full distillation training loop.

    Args:
        train_dataloader: Dataloader containing training dataset.
        val_dataloader: Optional dataloader for epoch-end validation.
        epochs: Total number of epochs to train up to (1-indexed, inclusive).
            Defaults to 10.
        callbacks: Optional list of callback functions triggered each epoch.
            A callback exposing a truthy `stop` attribute after being called
            (e.g. `EarlyStopping`) interrupts training at the end of that epoch.
        start_epoch: 1-based epoch index to resume training from. Defaults to 1
            (a fresh run). Used together with `history` when resuming from a
            checkpoint saved via `Distiller.save_checkpoint`.
        history: Existing training history to append to, as returned by a
            previous call to `fit`. Defaults to None (starts a fresh history).

    Returns:
        dict[str, list[float]]: Training history tracking loss and metrics,
        covering both the resumed epochs (if any) and the new ones.
    """
    if history is None:
        history = {
            "train_loss": [],
            "train_accuracy": [],
            "val_loss": [],
            "val_accuracy": [],
        }

    for epoch in range(start_epoch, epochs + 1):
        train_metrics = self.train_epoch(train_dataloader, epoch, epochs)
        history["train_loss"].append(train_metrics["loss"])
        history["train_accuracy"].append(train_metrics["accuracy"])

        val_metrics: dict[str, float] = {}
        if val_dataloader is not None:
            val_metrics = self.evaluate(val_dataloader)
            history["val_loss"].append(val_metrics["val_loss"])
            history["val_accuracy"].append(val_metrics["val_accuracy"])

        if self.scheduler is not None:
            self.scheduler.step()

        status = (
            f"Epoch [{epoch:02d}/{epochs:02d}] "
            f"Train Loss: {train_metrics['loss']:.4f} - "
            f"Train Acc: {train_metrics['accuracy']:.2f}%"
        )
        if val_metrics:
            status += (
                f" | Val Loss: {val_metrics['val_loss']:.4f} - "
                f"Val Acc: {val_metrics['val_accuracy']:.2f}%"
            )
        tqdm.write(status)

        if callbacks:
            epoch_summary = {**train_metrics, **val_metrics}
            for callback in callbacks:
                callback(epoch, epoch_summary)

            if any(getattr(callback, "stop", False) for callback in callbacks):
                tqdm.write(f"Training stopped early at epoch {epoch}/{epochs}.")
                break

    return history

train_epoch

train_epoch(
    dataloader: DataLoader,
    epoch_idx: int,
    total_epochs: int,
) -> dict[str, float]

Runs a single training epoch over the provided dataloader.

Parameters:

Name	Type	Description	Default
dataloader	DataLoader	Training dataloader yielding (inputs, labels) batches.	required
epoch_idx	int	Current 1-based epoch index.	required
total_epochs	int	Total number of planned epochs.	required

Returns:

Type	Description
dict[str, float]	dict[str, float]: Aggregated training metrics (loss, accuracy).

Source code in src/shrinkai/distillation/engine.py

def train_epoch(
    self,
    dataloader: DataLoader,
    epoch_idx: int,
    total_epochs: int,
) -> dict[str, float]:
    """Runs a single training epoch over the provided dataloader.

    Args:
        dataloader: Training dataloader yielding (inputs, labels) batches.
        epoch_idx: Current 1-based epoch index.
        total_epochs: Total number of planned epochs.

    Returns:
        dict[str, float]: Aggregated training metrics (loss, accuracy).
    """
    self.student.train()
    total_loss = 0.0
    correct = 0
    total_samples = 0

    desc = f"Epoch [{epoch_idx}/{total_epochs}] Training"
    pbar = tqdm(dataloader, desc=desc, leave=False)

    for batch in pbar:
        inputs, labels = self._unpack_batch(batch)

        with (
            torch.no_grad(),
            torch.autocast(
                device_type=self.device.type, dtype=self._amp_dtype, enabled=self.use_amp
            ),
        ):
            teacher_outputs = self._forward_model(self.teacher, inputs)

        self.optimizer.zero_grad()
        with torch.autocast(
            device_type=self.device.type, dtype=self._amp_dtype, enabled=self.use_amp
        ):
            student_outputs = self._forward_model(self.student, inputs)
            loss = self.criterion(
                student_outputs=student_outputs,
                teacher_outputs=teacher_outputs,
                labels=labels,
            )

        if not math.isfinite(loss.item()):
            raise RuntimeError(
                f"Loss diverged to {loss.item()} during training at epoch {epoch_idx}. "
                "Check learning rate or data scaling."
            )

        if self._use_scaler:
            self.scaler.scale(loss).backward()
            if self.grad_clip_norm is not None:
                self.scaler.unscale_(self.optimizer)
                nn.utils.clip_grad_norm_(self.student.parameters(), self.grad_clip_norm)
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            loss.backward()
            if self.grad_clip_norm is not None:
                nn.utils.clip_grad_norm_(self.student.parameters(), self.grad_clip_norm)
            self.optimizer.step()

        student_logits = (
            student_outputs[0] if isinstance(student_outputs, tuple) else student_outputs
        )

        if student_logits.dim() == 3:
            # LLM Causal Shift
            preds = torch.argmax(student_logits[..., :-1, :].contiguous(), dim=-1)
            shifted_labels = labels[..., 1:].contiguous()
            correct += (preds == shifted_labels).sum().item()
            batch_samples = shifted_labels.numel()
        else:
            preds = torch.argmax(student_logits, dim=-1)
            correct += (preds == labels).sum().item()
            batch_samples = labels.size(0)

        total_loss += loss.item() * batch_samples
        total_samples += batch_samples

        current_loss = total_loss / total_samples
        current_acc = (correct / total_samples) * 100.0
        pbar.set_postfix(loss=f"{current_loss:.4f}", acc=f"{current_acc:.2f}%")

    return {
        "loss": total_loss / total_samples,
        "accuracy": (correct / total_samples) * 100.0,
    }

Distiller

Unified, high-level facade for end-to-end knowledge distillation and benchmarking.

Examples:

>>> from shrinkai.distillation import Distiller
>>> distiller = Distiller(teacher=teacher_vit, student=student_mobilenet)
>>> distiller.fit(train_loader, val_loader, epochs=5)
>>> distiller.save_student("distilled_student.pt")

Methods:

Name	Description
__init__	Initializes the Distiller.
benchmark	Runs complete profiling suite on both models and outputs comparison report.
evaluate	Evaluates student performance on a given dataloader.
export_onnx	Exports the trained student to ONNX. See shrinkai.export.export_onnx.
export_torchscript	Exports the trained student to TorchScript. See
feature_analysis	Evaluates how well the student mimics the teacher's internal hidden states.
fit	Trains the student model using knowledge distillation.
load_checkpoint	Restores a full training checkpoint saved by save_checkpoint.
load_student	Loads trained weights into the student model.
save_checkpoint	Saves a full training checkpoint (student, optimizer, scheduler, history).
save_student	Saves trained student model weights to disk.

Source code in src/shrinkai/distillation/distiller.py

class Distiller:
    """Unified, high-level facade for end-to-end knowledge distillation and benchmarking.

    Examples:
        >>> from shrinkai.distillation import Distiller
        >>> distiller = Distiller(teacher=teacher_vit, student=student_mobilenet)
        >>> distiller.fit(train_loader, val_loader, epochs=5)
        >>> distiller.save_student("distilled_student.pt")
    """

    def __init__(
        self,
        teacher: nn.Module,
        student: nn.Module,
        criterion: BaseDistillationLoss | None = None,
        optimizer: torch.optim.Optimizer | Literal["adam", "adamw", "sgd"] = "adamw",
        lr: float = 1e-3,
        weight_decay: float = 1e-4,
        device: torch.device | str = "auto",
        scheduler: torch.optim.lr_scheduler._LRScheduler | None = None,
        use_amp: bool = False,
        grad_clip_norm: float | None = None,
        engine_class: type[DistillationEngine] | None = None,
    ) -> None:
        """Initializes the Distiller.

        Args:
            teacher: Pre-trained teacher neural network module.
            student: Target lightweight student neural network module to train.
            criterion: Distillation loss adhering to `BaseDistillationLoss`.
                Defaults to `HintonLoss()`.
            optimizer: PyTorch optimizer or name string ('adam', 'adamw', 'sgd').
                Defaults to 'adamw'.
            lr: Learning rate applied when creating default optimizer. Defaults to 1e-3.
            weight_decay: Weight decay factor for optimizer. Defaults to 1e-4.
            device: Computing target ('auto', 'mps', 'cuda', 'cpu' or torch.device).
                Defaults to 'auto'.
            scheduler: Optional learning rate scheduler updated per epoch.
            use_amp: If True, trains under mixed precision (fp16+scaling on CUDA,
                bf16 on CPU/MPS). Defaults to False.
            grad_clip_norm: If set, clips the student's gradient global L2 norm to
                this value before each optimizer step. Defaults to None.
            engine_class: The engine class for distillation.
                If None, it uses a basic DistillationEngine.
        """
        self.teacher = teacher
        self.student = student
        self.device = resolve_device(device)
        self.criterion = criterion if criterion is not None else HintonLoss()
        self.teacher.to(self.device)
        self.student.to(self.device)
        self.criterion.to(self.device)

        if isinstance(optimizer, str):
            self.optimizer = self._build_optimizer(
                opt_name=optimizer.lower(),
                lr=lr,
                weight_decay=weight_decay,
            )
        else:
            self.optimizer = optimizer

        self.scheduler = scheduler
        self.history: dict[str, list[float]] = {
            "train_loss": [],
            "train_accuracy": [],
            "val_loss": [],
            "val_accuracy": [],
        }
        if engine_class is None:
            self._engine = DistillationEngine(
                student=self.student,
                teacher=self.teacher,
                criterion=self.criterion,
                optimizer=self.optimizer,
                device=self.device,
                scheduler=self.scheduler,
                use_amp=use_amp,
                grad_clip_norm=grad_clip_norm,
            )
        else:
            self._engine = engine_class(
                student=self.student,
                teacher=self.teacher,
                criterion=self.criterion,
                optimizer=self.optimizer,
                device=self.device,
                scheduler=self.scheduler,
                use_amp=use_amp,
                grad_clip_norm=grad_clip_norm,
            )

    def _build_optimizer(
        self,
        opt_name: str,
        lr: float,
        weight_decay: float,
    ) -> torch.optim.Optimizer:
        """Builds standard optimizer for student parameters."""
        params = [p for p in self.student.parameters() if p.requires_grad]
        params += [p for p in self.criterion.parameters() if p.requires_grad]

        if opt_name == "adamw":
            return torch.optim.AdamW(params, lr=lr, weight_decay=weight_decay)
        if opt_name == "adam":
            return torch.optim.Adam(params, lr=lr, weight_decay=weight_decay)
        if opt_name == "sgd":
            return torch.optim.SGD(params, lr=lr, momentum=0.9, weight_decay=weight_decay)

        raise ValueError(
            f"Unsupported optimizer '{opt_name}'. Choose from 'adamw', 'adam', 'sgd' "
            "or pass an instantiated `torch.optim.Optimizer`."
        )

    def fit(
        self,
        train_dataloader: DataLoader,
        val_dataloader: DataLoader | None = None,
        epochs: int = 10,
        callbacks: list[Callable[[int, dict[str, float]], None]] | None = None,
        resume: bool = False,
    ) -> dict[str, list[float]]:
        """Trains the student model using knowledge distillation.

        Args:
            train_dataloader: Dataloader yielding training batches.
            val_dataloader: Optional dataloader for evaluation after each epoch.
            epochs: Total number of epochs to train up to (1-indexed, inclusive).
                Defaults to 10.
            callbacks: Optional list of callback functions triggered at epoch end.
                A callback exposing a truthy `stop` attribute (e.g. `EarlyStopping`)
                interrupts training at the end of that epoch.
            resume: If True, continues training from `self.history` (populated by a
                previous `fit()` call or by `load_checkpoint()`) instead of starting
                a fresh run from epoch 1. Defaults to False.

        Returns:
            dict[str, list[float]]: Dictionary tracking history across epochs. Also
            stored on `self.history` for later checkpointing.
        """
        start_epoch = len(self.history["train_loss"]) + 1 if resume else 1
        history = self.history if resume else None

        self.history = self._engine.fit(
            train_dataloader=train_dataloader,
            val_dataloader=val_dataloader,
            epochs=epochs,
            callbacks=callbacks,
            start_epoch=start_epoch,
            history=history,
        )
        return self.history

    def evaluate(self, dataloader: DataLoader) -> dict[str, float]:
        """Evaluates student performance on a given dataloader.

        Args:
            dataloader: Dataloader yielding validation/testing batches.

        Returns:
            dict[str, float]: Validation metrics dictionary.
        """
        return self._engine.evaluate(dataloader)

    def benchmark(
        self,
        sample_input: torch.Tensor,
        teacher_name: str = "Teacher",
        student_name: str = "Student",
        val_dataloader: DataLoader | None = None,
        compute_flops: bool = False,
    ) -> BenchmarkReport:
        """Runs complete profiling suite on both models and outputs comparison report.

        Args:
            sample_input: Batch tensor matching target inference dimension (e.g., [1, 3, 224, 224]).
            teacher_name: Display label for teacher model. Defaults to "Teacher".
            student_name: Display label for student model. Defaults to "Student".
            val_dataloader: Optional dataloader to compute final accuracy metrics.
            compute_flops: If True, also reports FLOPs per sample for both models.
                Defaults to False. See `Profiler.compare`.

        Returns:
            BenchmarkReport: Structured benchmark report ready for `.show()`.
        """
        teacher_acc: float | None = None
        student_acc: float | None = None

        if val_dataloader is not None:
            teacher_acc = compute_accuracy(self.teacher, val_dataloader, self.device)
            student_acc = compute_accuracy(self.student, val_dataloader, self.device)

        return Profiler.compare(
            teacher=self.teacher,
            student=self.student,
            sample_input=sample_input,
            device=self.device,
            teacher_name=teacher_name,
            student_name=student_name,
            teacher_acc=teacher_acc,
            student_acc=student_acc,
            compute_flops=compute_flops,
        )

    def feature_analysis(
        self, dataloader: DataLoader, metrics: list[str] | None = None
    ) -> FeatureAnalyzerReport:
        """Evaluates how well the student mimics the teacher's internal hidden states.

        This requires both the teacher and student models to have been wrapped
        with `FeatureExtractor` prior to initializing the Distiller.

        Args:
            dataloader: Dataloader yielding validation/testing batches.
            metrics: List of metrics to compute (e.g., 'cka'). Defaults to ["cka"].

        Returns:
            FeatureAnalyzerReport: Structured report ready for `.show()`.

        Raises:
            ValueError: If models are not wrapped with `FeatureExtractor`.
        """

        if not isinstance(self.teacher, FeatureExtractor) or not isinstance(
            self.student, FeatureExtractor
        ):
            raise ValueError(
                "Feature analysis requires both models to be wrapped with `FeatureExtractor`. "
                "If you are only doing logit-based distillation (HintonLoss, etc), internal states "
                "are not accessible."
            )

        if metrics is None:
            metrics = ["cka"]

        analyzer = FeatureAnalyzer(
            teacher_extractor=self.teacher,
            student_extractor=self.student,
            device=self.device,
        )

        return analyzer.evaluate(dataloader=dataloader, metrics=metrics)

    def save_student(self, path: str | Path) -> None:
        """Saves trained student model weights to disk.

        Args:
            path: Destination file path (e.g. 'models/student.pt').
        """
        save_path = Path(path)
        save_path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(self.student.state_dict(), save_path)

    def load_student(self, path: str | Path) -> None:
        """Loads trained weights into the student model.

        Args:
            path: File path of saved state dictionary.
        """
        load_path = Path(path)
        state_dict = torch.load(load_path, map_location=self.device, weights_only=True)
        self.student.load_state_dict(state_dict)

    def export_onnx(
        self,
        path: str | Path,
        sample_input: torch.Tensor | tuple[torch.Tensor, ...],
        **kwargs: Any,
    ) -> Path:
        """Exports the trained student to ONNX. See `shrinkai.export.export_onnx`.

        Args:
            path: Destination `.onnx` file path.
            sample_input: Representative input tensor (or tuple of tensors).
            **kwargs: Forwarded to `shrinkai.export.export_onnx` (e.g.
                `dynamic_batch`, `opset_version`, `input_names`).

        Returns:
            Path: The path the model was exported to.
        """
        return export_onnx(self.student, sample_input, path, **kwargs)

    def export_torchscript(
        self,
        path: str | Path,
        sample_input: torch.Tensor | tuple[torch.Tensor, ...] | None = None,
        method: Literal["trace", "script"] = "trace",
    ) -> Path:
        """Exports the trained student to TorchScript. See
        `shrinkai.export.export_torchscript`.

        Args:
            path: Destination file path.
            sample_input: Required when `method="trace"`.
            method: "trace" (default) or "script".

        Returns:
            Path: The path the model was exported to.
        """
        return export_torchscript(self.student, path, sample_input=sample_input, method=method)

    def save_checkpoint(self, path: str | Path) -> None:
        """Saves a full training checkpoint (student, optimizer, scheduler, history).

        Unlike `save_student`, which only persists inference weights, this saves
        everything needed to resume training later via `load_checkpoint` followed by
        `fit(..., resume=True)`.

        Args:
            path: Destination file path (e.g. 'checkpoints/epoch_10.pt').
        """
        save_path = Path(path)
        save_path.parent.mkdir(parents=True, exist_ok=True)
        checkpoint = {
            "student_state_dict": self.student.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "scheduler_state_dict": (
                self.scheduler.state_dict() if self.scheduler is not None else None
            ),
            "history": self.history,
        }
        torch.save(checkpoint, save_path)

    def load_checkpoint(self, path: str | Path) -> None:
        """Restores a full training checkpoint saved by `save_checkpoint`.

        After calling this, resume training with `fit(..., resume=True)`, it will
        continue from the epoch right after the last one recorded in the restored
        history.

        Args:
            path: File path of the saved checkpoint.

        Raises:
            ValueError: If this Distiller's scheduler configuration (present vs.
                absent) does not match the one the checkpoint was saved with.
        """
        load_path = Path(path)
        checkpoint = torch.load(load_path, map_location=self.device, weights_only=True)

        self.student.load_state_dict(checkpoint["student_state_dict"])
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

        scheduler_state = checkpoint.get("scheduler_state_dict")
        if self.scheduler is None and scheduler_state is not None:
            raise ValueError(
                "The checkpoint contains a scheduler state, but this Distiller "
                "has no scheduler configured."
            )
        if self.scheduler is not None and scheduler_state is None:
            raise ValueError(
                "This Distiller has a scheduler configured, but the checkpoint "
                "was saved without one."
            )
        if self.scheduler is not None:
            self.scheduler.load_state_dict(scheduler_state)

        self.history = checkpoint.get(
            "history",
            {"train_loss": [], "train_accuracy": [], "val_loss": [], "val_accuracy": []},
        )

Methods:

__init__

__init__(
    teacher: Module,
    student: Module,
    criterion: BaseDistillationLoss | None = None,
    optimizer: Optimizer
    | Literal["adam", "adamw", "sgd"] = "adamw",
    lr: float = 0.001,
    weight_decay: float = 0.0001,
    device: device | str = "auto",
    scheduler: _LRScheduler | None = None,
    use_amp: bool = False,
    grad_clip_norm: float | None = None,
    engine_class: type[DistillationEngine] | None = None,
) -> None

Initializes the Distiller.

Parameters:

Name	Type	Description	Default
teacher	Module	Pre-trained teacher neural network module.	required
student	Module	Target lightweight student neural network module to train.	required
criterion	BaseDistillationLoss | None	Distillation loss adhering to BaseDistillationLoss. Defaults to HintonLoss().	None
optimizer	Optimizer | Literal['adam', 'adamw', 'sgd']	PyTorch optimizer or name string ('adam', 'adamw', 'sgd'). Defaults to 'adamw'.	'adamw'
lr	float	Learning rate applied when creating default optimizer. Defaults to 1e-3.	0.001
weight_decay	float	Weight decay factor for optimizer. Defaults to 1e-4.	0.0001
device	device | str	Computing target ('auto', 'mps', 'cuda', 'cpu' or torch.device). Defaults to 'auto'.	'auto'
scheduler	_LRScheduler | None	Optional learning rate scheduler updated per epoch.	None
use_amp	bool	If True, trains under mixed precision (fp16+scaling on CUDA, bf16 on CPU/MPS). Defaults to False.	False
grad_clip_norm	float | None	If set, clips the student's gradient global L2 norm to this value before each optimizer step. Defaults to None.	None
engine_class	type[DistillationEngine] | None	The engine class for distillation. If None, it uses a basic DistillationEngine.	None

Source code in src/shrinkai/distillation/distiller.py

def __init__(
    self,
    teacher: nn.Module,
    student: nn.Module,
    criterion: BaseDistillationLoss | None = None,
    optimizer: torch.optim.Optimizer | Literal["adam", "adamw", "sgd"] = "adamw",
    lr: float = 1e-3,
    weight_decay: float = 1e-4,
    device: torch.device | str = "auto",
    scheduler: torch.optim.lr_scheduler._LRScheduler | None = None,
    use_amp: bool = False,
    grad_clip_norm: float | None = None,
    engine_class: type[DistillationEngine] | None = None,
) -> None:
    """Initializes the Distiller.

    Args:
        teacher: Pre-trained teacher neural network module.
        student: Target lightweight student neural network module to train.
        criterion: Distillation loss adhering to `BaseDistillationLoss`.
            Defaults to `HintonLoss()`.
        optimizer: PyTorch optimizer or name string ('adam', 'adamw', 'sgd').
            Defaults to 'adamw'.
        lr: Learning rate applied when creating default optimizer. Defaults to 1e-3.
        weight_decay: Weight decay factor for optimizer. Defaults to 1e-4.
        device: Computing target ('auto', 'mps', 'cuda', 'cpu' or torch.device).
            Defaults to 'auto'.
        scheduler: Optional learning rate scheduler updated per epoch.
        use_amp: If True, trains under mixed precision (fp16+scaling on CUDA,
            bf16 on CPU/MPS). Defaults to False.
        grad_clip_norm: If set, clips the student's gradient global L2 norm to
            this value before each optimizer step. Defaults to None.
        engine_class: The engine class for distillation.
            If None, it uses a basic DistillationEngine.
    """
    self.teacher = teacher
    self.student = student
    self.device = resolve_device(device)
    self.criterion = criterion if criterion is not None else HintonLoss()
    self.teacher.to(self.device)
    self.student.to(self.device)
    self.criterion.to(self.device)

    if isinstance(optimizer, str):
        self.optimizer = self._build_optimizer(
            opt_name=optimizer.lower(),
            lr=lr,
            weight_decay=weight_decay,
        )
    else:
        self.optimizer = optimizer

    self.scheduler = scheduler
    self.history: dict[str, list[float]] = {
        "train_loss": [],
        "train_accuracy": [],
        "val_loss": [],
        "val_accuracy": [],
    }
    if engine_class is None:
        self._engine = DistillationEngine(
            student=self.student,
            teacher=self.teacher,
            criterion=self.criterion,
            optimizer=self.optimizer,
            device=self.device,
            scheduler=self.scheduler,
            use_amp=use_amp,
            grad_clip_norm=grad_clip_norm,
        )
    else:
        self._engine = engine_class(
            student=self.student,
            teacher=self.teacher,
            criterion=self.criterion,
            optimizer=self.optimizer,
            device=self.device,
            scheduler=self.scheduler,
            use_amp=use_amp,
            grad_clip_norm=grad_clip_norm,
        )

benchmark

benchmark(
    sample_input: Tensor,
    teacher_name: str = "Teacher",
    student_name: str = "Student",
    val_dataloader: DataLoader | None = None,
    compute_flops: bool = False,
) -> BenchmarkReport

Runs complete profiling suite on both models and outputs comparison report.

Parameters:

Name	Type	Description	Default
sample_input	Tensor	Batch tensor matching target inference dimension (e.g., [1, 3, 224, 224]).	required
teacher_name	str	Display label for teacher model. Defaults to "Teacher".	'Teacher'
student_name	str	Display label for student model. Defaults to "Student".	'Student'
val_dataloader	DataLoader | None	Optional dataloader to compute final accuracy metrics.	None
compute_flops	bool	If True, also reports FLOPs per sample for both models. Defaults to False. See Profiler.compare.	False

Returns:

Name	Type	Description
BenchmarkReport	BenchmarkReport	Structured benchmark report ready for .show().

Source code in src/shrinkai/distillation/distiller.py

def benchmark(
    self,
    sample_input: torch.Tensor,
    teacher_name: str = "Teacher",
    student_name: str = "Student",
    val_dataloader: DataLoader | None = None,
    compute_flops: bool = False,
) -> BenchmarkReport:
    """Runs complete profiling suite on both models and outputs comparison report.

    Args:
        sample_input: Batch tensor matching target inference dimension (e.g., [1, 3, 224, 224]).
        teacher_name: Display label for teacher model. Defaults to "Teacher".
        student_name: Display label for student model. Defaults to "Student".
        val_dataloader: Optional dataloader to compute final accuracy metrics.
        compute_flops: If True, also reports FLOPs per sample for both models.
            Defaults to False. See `Profiler.compare`.

    Returns:
        BenchmarkReport: Structured benchmark report ready for `.show()`.
    """
    teacher_acc: float | None = None
    student_acc: float | None = None

    if val_dataloader is not None:
        teacher_acc = compute_accuracy(self.teacher, val_dataloader, self.device)
        student_acc = compute_accuracy(self.student, val_dataloader, self.device)

    return Profiler.compare(
        teacher=self.teacher,
        student=self.student,
        sample_input=sample_input,
        device=self.device,
        teacher_name=teacher_name,
        student_name=student_name,
        teacher_acc=teacher_acc,
        student_acc=student_acc,
        compute_flops=compute_flops,
    )

evaluate

evaluate(dataloader: DataLoader) -> dict[str, float]

Evaluates student performance on a given dataloader.

Parameters:

Name	Type	Description	Default
dataloader	DataLoader	Dataloader yielding validation/testing batches.	required

Returns:

Type	Description
dict[str, float]	dict[str, float]: Validation metrics dictionary.

Source code in src/shrinkai/distillation/distiller.py

def evaluate(self, dataloader: DataLoader) -> dict[str, float]:
    """Evaluates student performance on a given dataloader.

    Args:
        dataloader: Dataloader yielding validation/testing batches.

    Returns:
        dict[str, float]: Validation metrics dictionary.
    """
    return self._engine.evaluate(dataloader)

export_onnx

export_onnx(
    path: str | Path,
    sample_input: Tensor | tuple[Tensor, ...],
    **kwargs: Any,
) -> Path

Exports the trained student to ONNX. See shrinkai.export.export_onnx.

Parameters:

Name	Type	Description	Default
path	str | Path	Destination .onnx file path.	required
sample_input	Tensor | tuple[Tensor, ...]	Representative input tensor (or tuple of tensors).	required
**kwargs	Any	Forwarded to shrinkai.export.export_onnx (e.g. dynamic_batch, opset_version, input_names).	{}

Returns:

Name	Type	Description
Path	Path	The path the model was exported to.

Source code in src/shrinkai/distillation/distiller.py

def export_onnx(
    self,
    path: str | Path,
    sample_input: torch.Tensor | tuple[torch.Tensor, ...],
    **kwargs: Any,
) -> Path:
    """Exports the trained student to ONNX. See `shrinkai.export.export_onnx`.

    Args:
        path: Destination `.onnx` file path.
        sample_input: Representative input tensor (or tuple of tensors).
        **kwargs: Forwarded to `shrinkai.export.export_onnx` (e.g.
            `dynamic_batch`, `opset_version`, `input_names`).

    Returns:
        Path: The path the model was exported to.
    """
    return export_onnx(self.student, sample_input, path, **kwargs)

export_torchscript

export_torchscript(
    path: str | Path,
    sample_input: Tensor | tuple[Tensor, ...] | None = None,
    method: Literal["trace", "script"] = "trace",
) -> Path

Exports the trained student to TorchScript. See shrinkai.export.export_torchscript.

Parameters:

Name	Type	Description	Default
path	str | Path	Destination file path.	required
sample_input	Tensor | tuple[Tensor, ...] | None	Required when method="trace".	None
method	Literal['trace', 'script']	"trace" (default) or "script".	'trace'

Returns:

Name	Type	Description
Path	Path	The path the model was exported to.

Source code in src/shrinkai/distillation/distiller.py

def export_torchscript(
    self,
    path: str | Path,
    sample_input: torch.Tensor | tuple[torch.Tensor, ...] | None = None,
    method: Literal["trace", "script"] = "trace",
) -> Path:
    """Exports the trained student to TorchScript. See
    `shrinkai.export.export_torchscript`.

    Args:
        path: Destination file path.
        sample_input: Required when `method="trace"`.
        method: "trace" (default) or "script".

    Returns:
        Path: The path the model was exported to.
    """
    return export_torchscript(self.student, path, sample_input=sample_input, method=method)

feature_analysis

feature_analysis(
    dataloader: DataLoader, metrics: list[str] | None = None
) -> FeatureAnalyzerReport

Evaluates how well the student mimics the teacher's internal hidden states.

This requires both the teacher and student models to have been wrapped with FeatureExtractor prior to initializing the Distiller.

Parameters:

Name	Type	Description	Default
dataloader	DataLoader	Dataloader yielding validation/testing batches.	required
metrics	list[str] | None	List of metrics to compute (e.g., 'cka'). Defaults to ["cka"].	None

Returns:

Name	Type	Description
FeatureAnalyzerReport	FeatureAnalyzerReport	Structured report ready for .show().

Raises:

Type	Description
ValueError	If models are not wrapped with FeatureExtractor.

Source code in src/shrinkai/distillation/distiller.py

def feature_analysis(
    self, dataloader: DataLoader, metrics: list[str] | None = None
) -> FeatureAnalyzerReport:
    """Evaluates how well the student mimics the teacher's internal hidden states.

    This requires both the teacher and student models to have been wrapped
    with `FeatureExtractor` prior to initializing the Distiller.

    Args:
        dataloader: Dataloader yielding validation/testing batches.
        metrics: List of metrics to compute (e.g., 'cka'). Defaults to ["cka"].

    Returns:
        FeatureAnalyzerReport: Structured report ready for `.show()`.

    Raises:
        ValueError: If models are not wrapped with `FeatureExtractor`.
    """

    if not isinstance(self.teacher, FeatureExtractor) or not isinstance(
        self.student, FeatureExtractor
    ):
        raise ValueError(
            "Feature analysis requires both models to be wrapped with `FeatureExtractor`. "
            "If you are only doing logit-based distillation (HintonLoss, etc), internal states "
            "are not accessible."
        )

    if metrics is None:
        metrics = ["cka"]

    analyzer = FeatureAnalyzer(
        teacher_extractor=self.teacher,
        student_extractor=self.student,
        device=self.device,
    )

    return analyzer.evaluate(dataloader=dataloader, metrics=metrics)

fit

fit(
    train_dataloader: DataLoader,
    val_dataloader: DataLoader | None = None,
    epochs: int = 10,
    callbacks: list[Callable[[int, dict[str, float]], None]]
    | None = None,
    resume: bool = False,
) -> dict[str, list[float]]

Trains the student model using knowledge distillation.

Parameters:

Name	Type	Description	Default
train_dataloader	DataLoader	Dataloader yielding training batches.	required
val_dataloader	DataLoader | None	Optional dataloader for evaluation after each epoch.	None
epochs	int	Total number of epochs to train up to (1-indexed, inclusive). Defaults to 10.	10
callbacks	list[Callable[[int, dict[str, float]], None]] | None	Optional list of callback functions triggered at epoch end. A callback exposing a truthy stop attribute (e.g. EarlyStopping) interrupts training at the end of that epoch.	None
resume	bool	If True, continues training from self.history (populated by a previous fit() call or by load_checkpoint()) instead of starting a fresh run from epoch 1. Defaults to False.	False

Returns:

Type	Description
dict[str, list[float]]	dict[str, list[float]]: Dictionary tracking history across epochs. Also
dict[str, list[float]]	stored on self.history for later checkpointing.

Source code in src/shrinkai/distillation/distiller.py

def fit(
    self,
    train_dataloader: DataLoader,
    val_dataloader: DataLoader | None = None,
    epochs: int = 10,
    callbacks: list[Callable[[int, dict[str, float]], None]] | None = None,
    resume: bool = False,
) -> dict[str, list[float]]:
    """Trains the student model using knowledge distillation.

    Args:
        train_dataloader: Dataloader yielding training batches.
        val_dataloader: Optional dataloader for evaluation after each epoch.
        epochs: Total number of epochs to train up to (1-indexed, inclusive).
            Defaults to 10.
        callbacks: Optional list of callback functions triggered at epoch end.
            A callback exposing a truthy `stop` attribute (e.g. `EarlyStopping`)
            interrupts training at the end of that epoch.
        resume: If True, continues training from `self.history` (populated by a
            previous `fit()` call or by `load_checkpoint()`) instead of starting
            a fresh run from epoch 1. Defaults to False.

    Returns:
        dict[str, list[float]]: Dictionary tracking history across epochs. Also
        stored on `self.history` for later checkpointing.
    """
    start_epoch = len(self.history["train_loss"]) + 1 if resume else 1
    history = self.history if resume else None

    self.history = self._engine.fit(
        train_dataloader=train_dataloader,
        val_dataloader=val_dataloader,
        epochs=epochs,
        callbacks=callbacks,
        start_epoch=start_epoch,
        history=history,
    )
    return self.history

load_checkpoint

load_checkpoint(path: str | Path) -> None

Restores a full training checkpoint saved by save_checkpoint.

After calling this, resume training with fit(..., resume=True), it will continue from the epoch right after the last one recorded in the restored history.

Parameters:

Name	Type	Description	Default
path	str | Path	File path of the saved checkpoint.	required

Raises:

Type	Description
ValueError	If this Distiller's scheduler configuration (present vs. absent) does not match the one the checkpoint was saved with.

Source code in src/shrinkai/distillation/distiller.py

def load_checkpoint(self, path: str | Path) -> None:
    """Restores a full training checkpoint saved by `save_checkpoint`.

    After calling this, resume training with `fit(..., resume=True)`, it will
    continue from the epoch right after the last one recorded in the restored
    history.

    Args:
        path: File path of the saved checkpoint.

    Raises:
        ValueError: If this Distiller's scheduler configuration (present vs.
            absent) does not match the one the checkpoint was saved with.
    """
    load_path = Path(path)
    checkpoint = torch.load(load_path, map_location=self.device, weights_only=True)

    self.student.load_state_dict(checkpoint["student_state_dict"])
    self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

    scheduler_state = checkpoint.get("scheduler_state_dict")
    if self.scheduler is None and scheduler_state is not None:
        raise ValueError(
            "The checkpoint contains a scheduler state, but this Distiller "
            "has no scheduler configured."
        )
    if self.scheduler is not None and scheduler_state is None:
        raise ValueError(
            "This Distiller has a scheduler configured, but the checkpoint "
            "was saved without one."
        )
    if self.scheduler is not None:
        self.scheduler.load_state_dict(scheduler_state)

    self.history = checkpoint.get(
        "history",
        {"train_loss": [], "train_accuracy": [], "val_loss": [], "val_accuracy": []},
    )

load_student

load_student(path: str | Path) -> None

Loads trained weights into the student model.

Parameters:

Name	Type	Description	Default
path	str | Path	File path of saved state dictionary.	required

Source code in src/shrinkai/distillation/distiller.py

def load_student(self, path: str | Path) -> None:
    """Loads trained weights into the student model.

    Args:
        path: File path of saved state dictionary.
    """
    load_path = Path(path)
    state_dict = torch.load(load_path, map_location=self.device, weights_only=True)
    self.student.load_state_dict(state_dict)

save_checkpoint

save_checkpoint(path: str | Path) -> None

Saves a full training checkpoint (student, optimizer, scheduler, history).

Unlike save_student, which only persists inference weights, this saves everything needed to resume training later via load_checkpoint followed by fit(..., resume=True).

Parameters:

Name	Type	Description	Default
path	str | Path	Destination file path (e.g. 'checkpoints/epoch_10.pt').	required

Source code in src/shrinkai/distillation/distiller.py

def save_checkpoint(self, path: str | Path) -> None:
    """Saves a full training checkpoint (student, optimizer, scheduler, history).

    Unlike `save_student`, which only persists inference weights, this saves
    everything needed to resume training later via `load_checkpoint` followed by
    `fit(..., resume=True)`.

    Args:
        path: Destination file path (e.g. 'checkpoints/epoch_10.pt').
    """
    save_path = Path(path)
    save_path.parent.mkdir(parents=True, exist_ok=True)
    checkpoint = {
        "student_state_dict": self.student.state_dict(),
        "optimizer_state_dict": self.optimizer.state_dict(),
        "scheduler_state_dict": (
            self.scheduler.state_dict() if self.scheduler is not None else None
        ),
        "history": self.history,
    }
    torch.save(checkpoint, save_path)

save_student

save_student(path: str | Path) -> None

Saves trained student model weights to disk.

Parameters:

Name	Type	Description	Default
path	str | Path	Destination file path (e.g. 'models/student.pt').	required

Source code in src/shrinkai/distillation/distiller.py

def save_student(self, path: str | Path) -> None:
    """Saves trained student model weights to disk.

    Args:
        path: Destination file path (e.g. 'models/student.pt').
    """
    save_path = Path(path)
    save_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(self.student.state_dict(), save_path)

EarlyStopping

Stops training when a monitored metric has stopped improving.

Attributes:

Name	Type	Description
stop	bool	Set to True once patience is exhausted. Read by DistillationEngine.fit after each epoch to interrupt the loop early.

Methods:

Name	Description
__call__	Updates internal state and sets self.stop if patience is exhausted.
__init__	Initializes the EarlyStopping callback.

Source code in src/shrinkai/distillation/callbacks.py

class EarlyStopping:
    """Stops training when a monitored metric has stopped improving.

    Attributes:
        stop (bool): Set to True once `patience` is exhausted. Read by
            `DistillationEngine.fit` after each epoch to interrupt the loop early.
    """

    def __init__(
        self,
        monitor: str = "val_loss",
        patience: int = 5,
        mode: Literal["min", "max"] = "min",
        min_delta: float = 0.0,
    ) -> None:
        """Initializes the EarlyStopping callback.

        Args:
            monitor: Metric key to watch in the epoch summary dict (e.g. "val_loss").
            patience: Number of consecutive non-improving epochs tolerated before
                training is stopped.
            mode: "min" if lower values of `monitor` are better, "max" otherwise.
            min_delta: Minimum absolute change to qualify as an improvement.

        Raises:
            ValueError: If `mode` is not "min" or "max".
        """
        if mode not in ("min", "max"):
            raise ValueError(f"mode must be 'min' or 'max', got '{mode}'.")

        self.monitor = monitor
        self.patience = patience
        self.mode = mode
        self.min_delta = min_delta

        self.best_score: float | None = None
        self.num_bad_epochs = 0
        self.stop = False

    def _is_improvement(self, current: float) -> bool:
        if self.best_score is None:
            return True
        if self.mode == "min":
            return current < self.best_score - self.min_delta
        return current > self.best_score + self.min_delta

    def __call__(self, epoch: int, metrics: dict[str, float]) -> None:
        """Updates internal state and sets `self.stop` if patience is exhausted.

        Args:
            epoch: Current 1-based epoch index.
            metrics: Epoch summary dict, as passed by `DistillationEngine.fit`.
        """
        if self.monitor not in metrics:
            logger.warning(
                "EarlyStopping: metric '%s' not found in epoch %d summary. Skipping check.",
                self.monitor,
                epoch,
            )
            return

        current = metrics[self.monitor]
        if self._is_improvement(current):
            self.best_score = current
            self.num_bad_epochs = 0
        else:
            self.num_bad_epochs += 1

        if self.num_bad_epochs >= self.patience:
            logger.info(
                "EarlyStopping: '%s' did not improve for %d epoch(s). Stopping at epoch %d.",
                self.monitor,
                self.patience,
                epoch,
            )
            self.stop = True

Methods:

__call__

__call__(epoch: int, metrics: dict[str, float]) -> None

Updates internal state and sets self.stop if patience is exhausted.

Parameters:

Name	Type	Description	Default
epoch	int	Current 1-based epoch index.	required
metrics	dict[str, float]	Epoch summary dict, as passed by DistillationEngine.fit.	required

Source code in src/shrinkai/distillation/callbacks.py

def __call__(self, epoch: int, metrics: dict[str, float]) -> None:
    """Updates internal state and sets `self.stop` if patience is exhausted.

    Args:
        epoch: Current 1-based epoch index.
        metrics: Epoch summary dict, as passed by `DistillationEngine.fit`.
    """
    if self.monitor not in metrics:
        logger.warning(
            "EarlyStopping: metric '%s' not found in epoch %d summary. Skipping check.",
            self.monitor,
            epoch,
        )
        return

    current = metrics[self.monitor]
    if self._is_improvement(current):
        self.best_score = current
        self.num_bad_epochs = 0
    else:
        self.num_bad_epochs += 1

    if self.num_bad_epochs >= self.patience:
        logger.info(
            "EarlyStopping: '%s' did not improve for %d epoch(s). Stopping at epoch %d.",
            self.monitor,
            self.patience,
            epoch,
        )
        self.stop = True

__init__

__init__(
    monitor: str = "val_loss",
    patience: int = 5,
    mode: Literal["min", "max"] = "min",
    min_delta: float = 0.0,
) -> None

Initializes the EarlyStopping callback.

Parameters:

Name	Type	Description	Default
monitor	str	Metric key to watch in the epoch summary dict (e.g. "val_loss").	'val_loss'
patience	int	Number of consecutive non-improving epochs tolerated before training is stopped.	5
mode	Literal['min', 'max']	"min" if lower values of monitor are better, "max" otherwise.	'min'
min_delta	float	Minimum absolute change to qualify as an improvement.	0.0

Raises:

Type	Description
ValueError	If mode is not "min" or "max".

Source code in src/shrinkai/distillation/callbacks.py

def __init__(
    self,
    monitor: str = "val_loss",
    patience: int = 5,
    mode: Literal["min", "max"] = "min",
    min_delta: float = 0.0,
) -> None:
    """Initializes the EarlyStopping callback.

    Args:
        monitor: Metric key to watch in the epoch summary dict (e.g. "val_loss").
        patience: Number of consecutive non-improving epochs tolerated before
            training is stopped.
        mode: "min" if lower values of `monitor` are better, "max" otherwise.
        min_delta: Minimum absolute change to qualify as an improvement.

    Raises:
        ValueError: If `mode` is not "min" or "max".
    """
    if mode not in ("min", "max"):
        raise ValueError(f"mode must be 'min' or 'max', got '{mode}'.")

    self.monitor = monitor
    self.patience = patience
    self.mode = mode
    self.min_delta = min_delta

    self.best_score: float | None = None
    self.num_bad_epochs = 0
    self.stop = False

ModelCheckpoint

Saves a model's weights to disk during training.

Holds a direct reference to the module to save (typically the student), so it plugs into fit(callbacks=[...]) without changing the existing (epoch, metrics) -> None callback signature.

Methods:

Name	Description
__call__	Saves the model's weights, respecting save_best_only.
__init__	Initializes the ModelCheckpoint callback.

Source code in src/shrinkai/distillation/callbacks.py

class ModelCheckpoint:
    """Saves a model's weights to disk during training.

    Holds a direct reference to the module to save (typically the student), so it
    plugs into `fit(callbacks=[...])` without changing the existing
    `(epoch, metrics) -> None` callback signature.
    """

    def __init__(
        self,
        model: nn.Module,
        filepath: str | Path,
        monitor: str = "val_loss",
        mode: Literal["min", "max"] = "min",
        save_best_only: bool = True,
    ) -> None:
        """Initializes the ModelCheckpoint callback.

        Args:
            model: The module whose `state_dict()` is saved (e.g. `distiller.student`).
            filepath: Destination path for the saved weights.
            monitor: Metric key to watch when `save_best_only` is True.
            mode: "min" if lower values of `monitor` are better, "max" otherwise.
            save_best_only: If True, only overwrite `filepath` when `monitor` improves
                over its best value so far. If False, save unconditionally every epoch.

        Raises:
            ValueError: If `mode` is not "min" or "max".
        """
        if mode not in ("min", "max"):
            raise ValueError(f"mode must be 'min' or 'max', got '{mode}'.")

        self.model = model
        self.filepath = Path(filepath)
        self.monitor = monitor
        self.mode = mode
        self.save_best_only = save_best_only
        self.best_score: float | None = None

    def _is_improvement(self, current: float) -> bool:
        if self.best_score is None:
            return True
        if self.mode == "min":
            return current < self.best_score
        return current > self.best_score

    def _save(self) -> None:
        self.filepath.parent.mkdir(parents=True, exist_ok=True)
        torch.save(self.model.state_dict(), self.filepath)

    def __call__(self, epoch: int, metrics: dict[str, float]) -> None:
        """Saves the model's weights, respecting `save_best_only`.

        Args:
            epoch: Current 1-based epoch index.
            metrics: Epoch summary dict, as passed by `DistillationEngine.fit`.
        """
        if not self.save_best_only:
            self._save()
            return

        if self.monitor not in metrics:
            logger.warning(
                "ModelCheckpoint: metric '%s' not found in epoch %d summary. Skipping save.",
                self.monitor,
                epoch,
            )
            return

        current = metrics[self.monitor]
        if self._is_improvement(current):
            self.best_score = current
            self._save()
            logger.info(
                "ModelCheckpoint: '%s' improved to %.4f at epoch %d. Saved to %s.",
                self.monitor,
                current,
                epoch,
                self.filepath,
            )

Methods:

__call__

__call__(epoch: int, metrics: dict[str, float]) -> None

Saves the model's weights, respecting save_best_only.

Parameters:

Name	Type	Description	Default
epoch	int	Current 1-based epoch index.	required
metrics	dict[str, float]	Epoch summary dict, as passed by DistillationEngine.fit.	required

Source code in src/shrinkai/distillation/callbacks.py

def __call__(self, epoch: int, metrics: dict[str, float]) -> None:
    """Saves the model's weights, respecting `save_best_only`.

    Args:
        epoch: Current 1-based epoch index.
        metrics: Epoch summary dict, as passed by `DistillationEngine.fit`.
    """
    if not self.save_best_only:
        self._save()
        return

    if self.monitor not in metrics:
        logger.warning(
            "ModelCheckpoint: metric '%s' not found in epoch %d summary. Skipping save.",
            self.monitor,
            epoch,
        )
        return

    current = metrics[self.monitor]
    if self._is_improvement(current):
        self.best_score = current
        self._save()
        logger.info(
            "ModelCheckpoint: '%s' improved to %.4f at epoch %d. Saved to %s.",
            self.monitor,
            current,
            epoch,
            self.filepath,
        )

__init__

__init__(
    model: Module,
    filepath: str | Path,
    monitor: str = "val_loss",
    mode: Literal["min", "max"] = "min",
    save_best_only: bool = True,
) -> None

Initializes the ModelCheckpoint callback.

Parameters:

Name	Type	Description	Default
model	Module	The module whose state_dict() is saved (e.g. distiller.student).	required
filepath	str | Path	Destination path for the saved weights.	required
monitor	str	Metric key to watch when save_best_only is True.	'val_loss'
mode	Literal['min', 'max']	"min" if lower values of monitor are better, "max" otherwise.	'min'
save_best_only	bool	If True, only overwrite filepath when monitor improves over its best value so far. If False, save unconditionally every epoch.	True

Raises:

Type	Description
ValueError	If mode is not "min" or "max".

Source code in src/shrinkai/distillation/callbacks.py

def __init__(
    self,
    model: nn.Module,
    filepath: str | Path,
    monitor: str = "val_loss",
    mode: Literal["min", "max"] = "min",
    save_best_only: bool = True,
) -> None:
    """Initializes the ModelCheckpoint callback.

    Args:
        model: The module whose `state_dict()` is saved (e.g. `distiller.student`).
        filepath: Destination path for the saved weights.
        monitor: Metric key to watch when `save_best_only` is True.
        mode: "min" if lower values of `monitor` are better, "max" otherwise.
        save_best_only: If True, only overwrite `filepath` when `monitor` improves
            over its best value so far. If False, save unconditionally every epoch.

    Raises:
        ValueError: If `mode` is not "min" or "max".
    """
    if mode not in ("min", "max"):
        raise ValueError(f"mode must be 'min' or 'max', got '{mode}'.")

    self.model = model
    self.filepath = Path(filepath)
    self.monitor = monitor
    self.mode = mode
    self.save_best_only = save_best_only
    self.best_score: float | None = None