A Quick Start¶
At the end of this quick start, you will have a clear and complete overview of ShrinkAI, but you will also understand how to use it in order to distill and compress pyTorch neural networks, making them smaller and faster.
Installation¶
pip install shrinkai
Distillation¶
Distillation is a deep learning technique that allows a small model, called the student, to learn from a bigger and more complex model, the teacher. It often leads to better training, due to the information given by the teacher, such as logits or hidden states, that are more meaningful than raw data. Indeed, the teacher already learned and understood the hidden patterns from the data, so the student just has to copy them.
The module shrinkai.distillation provides everything needed to build easy distillation pipeline. It relies on a Distiller, to which you must provide both models, a loss function, and training parameters. The core of a distillation lies in the choice of the loss. The module shrinkai.distillation.losses provides many losses used for distillation, which all inherit from BaseDistillationLoss. The most famous one is HintonLoss, but you can check the API reference to see the list of them, choose the one that fit your problem, or even create your own.
from shrinkai.distillation import Distiller
from shrinkai.distillation.losses import HintonLoss
In the following, let's assume both teacher and student models are loaded, as well as necessary data loaders and a sample_input tensor matching the model's expected input shape (used below for pruning and export).
hinton_loss = HintonLoss()
distiller = Distiller(teacher=teacher, student=student, criterion=hinton_loss, optimizer="adamw")
distiller.fit(
train_dataloader=train_loader,
epochs=10,
)
distiller.save_student("your/model/path.pt")
With the code above, you are now able to create your own Distiller to train the student model and save it! Be aware that the package provides many losses and that the distiller has other methods, especially for evaluation, such as benchmark and feature_analysis. As it is just a quick start, these methods won't be introduced, but you can refer to the tutorials and the API reference. Tutorials 01, 02, and 03 will focus on knowledge distillation.
Compression¶
Once trained, a model can be shrunk further with shrinkai.compression: Pruner zeroes out redundant weights (mask-based, works on any architecture), ChannelPruner physically removes entire channels for real size/latency gains, and Quantizer reduces numerical precision (e.g. to int8). See Tutorial 04 for the full picture, including which to reach for and when.
from shrinkai.compression.pruning import Pruner, PruningConfig
pruner = Pruner(PruningConfig(method="unstructured", amount=0.3))
pruner.apply(distiller.student) # attaches masks in place
Pruner.finalize(distiller.student) # bakes them into the weights, removes the masking hooks
End-to-end pipeline¶
Putting it all together, a typical pipeline distills a student, compresses it, then exports it for deployment:
distiller = Distiller(teacher=teacher, student=student, criterion=HintonLoss())
distiller.fit(train_dataloader=train_loader, epochs=10)
Pruner(PruningConfig(amount=0.3)).apply(distiller.student)
Pruner.finalize(distiller.student)
distiller.export_onnx("model.onnx", sample_input)
Export Model¶
A trained (and optionally compressed) model can be exported out of the PyTorch process entirely: export_onnx targets ONNX Runtime/TensorRT/CoreML conversion pipelines, export_torchscript targets LibTorch (C++) or PyTorch Mobile. Both are available directly off Distiller (or standalone, from shrinkai.export, for any plain nn.Module). See Tutorial 06 for the full picture, including what does and doesn't work with pruned/quantized models.
distiller.export_onnx("model.onnx", sample_input) # requires `pip install shrinkai[export]`
distiller.export_torchscript("model.pt", sample_input)
Tutorials' content¶
In addition to this quick start notebook, it is recommended to read the tutorials. They provide a more detailed overview, with more explanations (both theoretical and practical) about the package. They were created in a specific order: topics covered in the earlier tutorials are later touched on briefly or not covered at all. Even if they cover most of the package's features, they do not cover all of them. For example, they will not show how to distill LLM for a next token prediction task, but it could be done using the appropriate losses.
Here is the tutorials' content:
- 01: Distillation of Vision models on CIFAR10 —
HintonLoss,FeatureLoss, combining losses, feature alignment metrics (CKA/RSA/Spatial Attention). - 02: Distillation of Transformers on SST-2 — adapting HuggingFace models to
shrinkai,AttentionMapLoss. - 03: Distillation of Language Models for text generation —
ReverseKLLoss, handling mismatched tokenizers/vocabularies. - 04: Model Compression (Pruning & Quantization) —
Pruner,ChannelPruner,Quantizer(PTQ/QAT), combining compression with distillation. - 05: Customization & Full Training Control — custom losses, callbacks (
EarlyStopping,ModelCheckpoint), mixed precision, gradient clipping, checkpoint/resume, custom optimizers, and fully custom training loops viaDistillationEngine. - 06: Exporting Models for Deployment —
export_onnx,export_torchscript, and how they interact with pruned/quantized models.