ShrinkAI: knowledge distillation and compression package

projects
1

I just released shrinkai, a python package for reducing neural network size and inference latency (distillation + pruning + quantization), that could be useful for edge/mobile deployment. Built to feel native if you already know PyTorch, there is no custom training loop required, except if you need one. Here are the main objects:

  • Distiller, allowing distillation from a teacher to a student, with already provided losses (logit-based, feature-based, attention transfer, hybrid, …)
  • Pruner / ChannelPruner, magnitude-based masking or physical channel pruning
  • Quantizer,PTQ and QAT
  • Profiler, params / size / latency / accuracy benchmarking, teacher vs. student

Here is a little snippet of code for distillation:

from shrinkai.distillation import Distiller
from shrinkai.distillation.losses import HintonLoss

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")

I tested on many examples you can see on the docs, but here are the result from the CIFAR-10 tutorial in just a few lines of code.

Config:

  • Teacher: VGG19-BN
  • Student: ResNet-20
  • Loss: HybridLoss that aggregates HintonLoss and FeatureLoss

Results:

  • 72.7x smaller (78.53 MB → 1.08 MB)
  • 5.3x faster inference
  • 92.6% of teacher accuracy retained

Feel free to use it for your own experiments and give feedbacks!

Install: pip install shrinkai
GitHub: https://github.com/elouanzer/shrinkai
Docs: ShrinkAI

Built this as a solo project because I noticed people kept rewriting the same distillation/compression code for different models. Still early (v0.1.0), feedback, issues, and contributions are all welcome. :slight_smile: