05: Fully Controled Training¶
The other tutorials focus on what to distill (i.e. a loss, a compression technique). This one focuses on how to run training itself, once a basic Distiller.fit() call is not quite enough anymore. It covers:
- Writing a custom distillation loss
- Callbacks with
EarlyStoppingandModelCheckpoint - Mixed precision and gradient clipping
- Checkpointing and resuming a training run across separate
Distillerinstances - Using your own Optimizer
- Writing a fully custom training loop
The running example reuses the same setup as in Tutorial 01 (a vgg19_bn teacher distilling into an untrained resnet20 student on CIFAR-10) purely as a vehicle. The point here is the training machinery around it, not the accuracy reached, so subsets and epoch counts are kept small on purpose.
Setup¶
import tempfile
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms
from shrinkai.distillation import Distiller, DistillationEngine, EarlyStopping, ModelCheckpoint
from shrinkai.distillation.losses import BaseDistillationLoss, HintonLoss
def build_cifar_dataloaders(batch_size: int = 32, train_size: int = 256, val_size: int = 128):
transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
]
)
train_dataset = Subset(
datasets.CIFAR10(root="./data", train=True, download=True, transform=transform),
range(train_size),
)
val_dataset = Subset(
datasets.CIFAR10(root="./data", train=False, download=True, transform=transform),
range(val_size),
)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
return train_loader, val_loader
def build_teacher_student():
teacher = torch.hub.load(
"chenyaofo/pytorch-cifar-models", "cifar10_vgg19_bn", pretrained=True, verbose=False
)
student = torch.hub.load(
"chenyaofo/pytorch-cifar-models", "cifar10_resnet20", pretrained=False, verbose=False
)
return teacher, student
train_loader, val_loader = build_cifar_dataloaders()
Writing a custom loss¶
Every loss in shrinkai.distillation.losses is just a subclass of BaseDistillationLoss, an nn.Module with one abstract method to implement:
def forward(self, student_outputs, teacher_outputs, labels=None) -> torch.Tensor: ...
Distiller/DistillationEngine have no special knowledge of the built-in losses, they only call criterion(student_outputs=..., teacher_outputs=..., labels=...). As an example, let's write a loss that is not in the library: instead of matching softened probability distributions (like HintonLoss), it pushes the student's raw logit vector to point in the same direction as the teacher's, via cosine similarity, blended with the standard cross-entropy on ground truth.
class CosineLogitLoss(BaseDistillationLoss):
"""Encourages the student's logits to align directionally with the
teacher's (cosine similarity), blended with cross-entropy on labels.
"""
def __init__(self, alpha: float = 0.5) -> None:
super().__init__()
if not (0.0 <= alpha <= 1.0):
raise ValueError(f"alpha must be between 0.0 and 1.0, got {alpha}")
self.alpha = alpha
self.ce = nn.CrossEntropyLoss()
def forward(self, student_outputs, teacher_outputs, labels=None):
# FeatureExtractor-wrapped models return (logits, features_dict) tuples;
# unwrap defensively so this loss works with or without one.
s_logits = student_outputs[0] if isinstance(student_outputs, tuple) else student_outputs
t_logits = teacher_outputs[0] if isinstance(teacher_outputs, tuple) else teacher_outputs
cosine_loss = 1.0 - F.cosine_similarity(s_logits, t_logits, dim=-1).mean()
if self.alpha == 1.0 or labels is None:
return cosine_loss
ce_loss = self.ce(s_logits, labels)
return (1.0 - self.alpha) * ce_loss + self.alpha * cosine_loss
It plugs into Distiller exactly like any built-in loss — nothing else changes:
teacher, student = build_teacher_student()
distiller = Distiller(
teacher=teacher,
student=student,
criterion=CosineLogitLoss(alpha=0.5),
optimizer="adamw",
lr=5e-4,
device="cpu",
)
_ = distiller.fit(train_loader, epochs=1)
Epoch [01/01] Train Loss: 1.6759 - Train Acc: 8.20%
Callbacks¶
The method fit(..., callbacks=[...]) accepts any list of callables matching (epoch: int, metrics: dict[str, float]) -> None, the same epoch summary dict (train_loss, train_accuracy, and val_loss/val_accuracy when a validation dataloader is given) is passed to every callback, every epoch. shrinkai.distillation ships two ready-made ones:
EarlyStopping(monitor, patience, mode, min_delta): after being called, it exposes astopattribute.fit()checksgetattr(callback, "stop", False)on every callback after each epoch and breaks out of the loop if any of them is truthy, a plain function callback is unaffected by this check, so it costs nothing to existing code.ModelCheckpoint(model, filepath, monitor, mode, save_best_only): holds a direct reference to the model to save (typicallydistiller.student), and writes it to disk whenevermonitorimproves.
Let's train for up to 10 epochs on our small subset, but stop as soon as validation loss has not improved for 2 epochs in a row, saving the best student found along the way:
teacher, student = build_teacher_student()
distiller = Distiller(
teacher=teacher,
student=student,
criterion=HintonLoss(temperature=4.0, alpha=0.7),
optimizer="adamw",
lr=5e-4,
device="cpu",
)
early_stopping = EarlyStopping(monitor="val_loss", patience=2, mode="min")
checkpoint_dir = tempfile.mkdtemp()
checkpoint = ModelCheckpoint(
distiller.student,
filepath=f"{checkpoint_dir}/best_student.pt",
monitor="val_loss",
mode="min",
)
history = distiller.fit(
train_loader,
val_loader,
epochs=10,
callbacks=[early_stopping, checkpoint],
)
print(f"Ran {len(history['train_loss'])} epochs (out of 10 requested).")
print(f"Stopped early: {early_stopping.stop} | best val_loss seen: {early_stopping.best_score:.4f}")
Epoch [01/10] Train Loss: 11.6079 - Train Acc: 12.89% | Val Loss: 10.5919 - Val Acc: 10.16%
Epoch [02/10] Train Loss: 10.7163 - Train Acc: 23.44% | Val Loss: 9.8844 - Val Acc: 21.88%
Epoch [03/10] Train Loss: 9.9858 - Train Acc: 31.64% | Val Loss: 9.4197 - Val Acc: 27.34%
Epoch [04/10] Train Loss: 9.4645 - Train Acc: 33.20% | Val Loss: 9.1767 - Val Acc: 28.91%
Epoch [05/10] Train Loss: 8.9022 - Train Acc: 36.72% | Val Loss: 8.8921 - Val Acc: 26.56%
Epoch [06/10] Train Loss: 8.5311 - Train Acc: 41.80% | Val Loss: 8.6197 - Val Acc: 28.91%
Epoch [07/10] Train Loss: 8.1278 - Train Acc: 43.75% | Val Loss: 8.6821 - Val Acc: 29.69%
Epoch [08/10] Train Loss: 7.8574 - Train Acc: 49.22% | Val Loss: 9.4021 - Val Acc: 24.22% Training stopped early at epoch 8/10. Ran 8 epochs (out of 10 requested). Stopped early: True | best val_loss seen: 8.6197
Training stopped well before the requested 10 epochs, the moment validation loss failed to improve for 2 epochs in a row, and checkpoint_dir now contains the student's weights from the best epoch, not the last one, ready to be reloaded with distiller.load_student(...) regardless of how training ended.
Mixed precision and gradient clipping¶
Two more constructor flags on Distiller (forwarded to DistillationEngine):
use_amp=True: runs the forward passes and loss computation undertorch.autocast. On CUDA this uses fp16 with aGradScaler(fp16 has too little dynamic range to skip loss-scaling safely); on CPU/MPS it uses bf16, which does not need a scaler (same exponent range as fp32).DistillationEnginepicks the right one automatically from the resolved device.grad_clip_norm=<float>: clips the student's gradient global L2 norm to that value right beforeoptimizer.step(), applied after unscaling when aGradScaleris active, a common stabilizer when logits (and therefore some losses' gradients) can occasionally spike.
Both compose with everything else shown in this notebook, callbacks, checkpointing, and custom losses are unaffected by either flag.
teacher, student = build_teacher_student()
distiller = Distiller(
teacher=teacher,
student=student,
criterion=HintonLoss(temperature=4.0, alpha=0.7),
optimizer="adamw",
lr=5e-4,
device="auto",
use_amp=True,
grad_clip_norm=1.0,
)
print(f"Resolved device: {distiller.device} | use_amp={distiller._engine.use_amp}")
_ = distiller.fit(train_loader, epochs=1)
Resolved device: mps | use_amp=True
Epoch [01/01] Train Loss: 11.7148 - Train Acc: 13.28%
Checkpointing and resuming¶
save_student()/load_student() (seen in Tutorial 01) only persist the student's weights, which is enough for inference but not for resuming training. save_checkpoint()/load_checkpoint() additionally capture the optimizer state, the scheduler state (if any), and the training history, so a run can be paused and resumed later, including from a brand new Distiller instance (e.g. after a process restart).
teacher, student = build_teacher_student()
distiller = Distiller(
teacher=teacher, student=student, criterion=HintonLoss(alpha=0.7),
optimizer="adamw", lr=5e-4, device="cpu",
)
_ = distiller.fit(train_loader, epochs=2)
checkpoint_path = f"{tempfile.mkdtemp()}/checkpoint.pt"
distiller.save_checkpoint(checkpoint_path)
print(f"Saved after {len(distiller.history['train_loss'])} epochs.")
Epoch [01/02] Train Loss: 11.5655 - Train Acc: 14.84%
Epoch [02/02] Train Loss: 10.8989 - Train Acc: 21.09% Saved after 2 epochs.
A different Distiller, same architecture, freshly constructed, can pick the run back up with load_checkpoint() followed by fit(..., resume=True). resume=True is what tells fit() to continue from len(self.history["train_loss"]) + 1 instead of restarting at epoch 1 (the default, resume=False, keeps the original behavior of every earlier tutorial unchanged).
new_teacher, new_student = build_teacher_student()
resumed_distiller = Distiller(
teacher=new_teacher, student=new_student, criterion=HintonLoss(alpha=0.7),
optimizer="adamw", lr=5e-4, device="cpu",
)
resumed_distiller.load_checkpoint(checkpoint_path)
print(f"Loaded checkpoint at epoch {len(resumed_distiller.history['train_loss'])}.")
resumed_distiller.fit(train_loader, epochs=4, resume=True)
print(f"Final history covers {len(resumed_distiller.history['train_loss'])} epochs.")
Loaded checkpoint at epoch 2.
Epoch [03/04] Train Loss: 10.3635 - Train Acc: 26.17%
Epoch [04/04] Train Loss: 9.7742 - Train Acc: 33.20% Final history covers 4 epochs.
Custom optimizer¶
Every example so far builds its optimizer from a string (optimizer="adamw" and lr/weight_decay parameters), as Distiller constructs a AdamW/Adam/SGD over the student's (and criterion's, if it has learnable parameters) parameters. For anything more specific, a different optimizer entirely, per-parameter-group learning rates, a custom weight decay schedule, or an optimizer from another library, you can simply build it yourself and pass the instance directly instead of a string.
teacher, student = build_teacher_student()
custom_optimizer = torch.optim.SGD(student.parameters(), lr=1e-2, momentum=0.9, nesterov=True)
distiller = Distiller(
teacher=teacher,
student=student,
criterion=HintonLoss(alpha=0.7),
optimizer=custom_optimizer, # an instance instead of a string
device="cpu",
)
_ = distiller.fit(train_loader, epochs=1)
Epoch [01/01] Train Loss: 11.5845 - Train Acc: 14.06%
Note: lr and weight_decay in Distiller are only used to build the optimizer when optimizer is a string. Once you pass an actual torch.optim.Optimizer instance, they are silently ignored, since it already carries its own hyperparameters. Keep in mind that Distiller's string-based path automatically includes the criterion's parameters in the optimizer (useful for losses like ProjectedFeatureLoss that hold a learnable projector), but with a custom optimizer instance, it is your own responsibility.
A fully custom training loop¶
Callbacks cover "run some code at the end of an epoch." Sometimes that is not enough, e.g. simulating a larger batch size than fits in memory via gradient accumulation, which needs to change when optimizer.step() is called, in the middle of the batch loop. That is exactly what Distiller delegates to under the hood: DistillationEngine, a public class on its own, meant to be subclassed for this kind of control. Overriding train_epoch (and/or evaluate, fit) gets you the full loop, with self.student, self.teacher, self.criterion, self.optimizer, self._unpack_batch, and self._forward_model all available exactly as the base class uses them.
class GradientAccumulationEngine(DistillationEngine):
"""Accumulates gradients over `accumulation_steps` batches before each
optimizer step, simulating a larger effective batch size.
"""
def __init__(self, *args, accumulation_steps: int = 4, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.accumulation_steps = accumulation_steps
def train_epoch(self, dataloader, epoch_idx, total_epochs):
self.student.train()
total_loss, correct, total_samples = 0.0, 0, 0
self.optimizer.zero_grad()
for step, batch in enumerate(dataloader):
inputs, labels = self._unpack_batch(batch)
with torch.no_grad():
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
)
(loss / self.accumulation_steps).backward()
if (step + 1) % self.accumulation_steps == 0:
self.optimizer.step()
self.optimizer.zero_grad()
student_logits = (
student_outputs[0] if isinstance(student_outputs, tuple) else student_outputs
)
preds = torch.argmax(student_logits, dim=-1)
batch_size = labels.size(0)
correct += (preds == labels).sum().item()
total_loss += loss.item() * batch_size
total_samples += batch_size
return {
"loss": total_loss / total_samples,
"accuracy": (correct / total_samples) * 100.0,
}
DistillationEngine (and therefore any subclass of it) is fully usable on its own, without going through Distiller at all — it just needs an optimizer built manually, since Distiller normally does that for you:
teacher, student = build_teacher_student()
distiller = Distiller(
teacher=teacher,
student=student,
criterion=HintonLoss(alpha=0.7),
optimizer="adamw",
lr=5e-4,
device="auto",
engine_class=GradientAccumulationEngine, # add your custom class here
)
_ = distiller.fit(train_loader, epochs=1)
Epoch [01/01] Train Loss: 11.5208 - Train Acc: 13.67%
Summary¶
| Need | Reach for |
|---|---|
A distillation objective not in shrinkai.distillation.losses |
Subclass BaseDistillationLoss |
| Stop training automatically / keep the best checkpoint | EarlyStopping, ModelCheckpoint in fit(callbacks=[...]) |
| Faster training / larger effective batches on limited memory | use_amp=True; gradient accumulation via a custom DistillationEngine |
| Exploding gradients | grad_clip_norm=<value> |
| Resuming a run later, possibly in a new process | save_checkpoint() / load_checkpoint() + fit(resume=True) |
| A different optimizer, per-group learning rates, or a custom schedule | Pass a torch.optim.Optimizer instance as optimizer= instead of a string |
| Changing when the optimizer steps, or the loop itself | Subclass DistillationEngine and override train_epoch/evaluate/fit |
None of these are mutually exclusive, a production training script would typically combine a custom loss, both callbacks, use_amp, and checkpointing in the same fit() call.
References¶
Prechelt, L. (1998). Early Stopping — But When?. In Neural Networks: Tricks of the Trade, LNCS 1524, Springer.
DOI: 10.1007/3-540-49430-8_3Micikevicius, P., Narang, S., Alben, J., et al. (2018). Mixed Precision Training. ICLR. arXiv:1710.03740.
https://arxiv.org/abs/1710.03740Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the Knowledge in a Neural Network. arXiv:1503.02531.
https://arxiv.org/abs/1503.02531