ShrinkAI Logo

Python versions License

ShrinkAI

ShrinkAI is a package for reducing neural networks size, making them ideal to run on small devices such as smartphones or robots, and for speeding up inference, particularly interesting for edge AI. This package includes, among other things, numerous distillation and compression techniques, all wrapped in an API that is easy to use for users familiar with PyTorch.

Installation

ShrinkAI is on pypi and can be installed with the following command:

pip install shrinkai

The package is compatible with Python 3.11+, and depends on torch, torchvision, rich, tqdm, and psutil. Exporting to ONNX additionally requires the export extra:

pip install shrinkai[export]

Quick Start

Distilling a smaller student model from a larger teacher only takes a Distiller and a loss:

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

# teacher, student: your PyTorch models. train_loader: your DataLoader.
distiller = Distiller(
    teacher=teacher,
    student=student,
    criterion=HintonLoss(),
    optimizer="adamw",
)

distiller.fit(train_dataloader=train_loader, epochs=10)
distiller.save_student("student.pt")

shrinkai also covers pruning, quantization, benchmarking, and exporting the trained model for deployment (ONNX, TorchScript). See the full documentation for the complete quick start, tutorials, and API reference.

Architecture

There is no flat top-level API: everything is imported from its submodule, e.g. from shrinkai.distillation import Distiller.

Module	What it provides
`shrinkai.distillation`	`Distiller` (facade), `DistillationEngine`, training callbacks, and a library of losses under `shrinkai.distillation.losses`
`shrinkai.adapters`	`FeatureExtractor` (hook-based intermediate activations) and dimension-matching projectors
`shrinkai.analysis`	Representation alignment metrics between teacher and student (CKA, RSA, Spatial Attention)
`shrinkai.compression`	`Pruner` / `ChannelPruner` (mask-based and physical pruning) and `Quantizer` (PTQ/QAT)
`shrinkai.profiler`	`Profiler`, `count_flops`, latency and memory measurement
`shrinkai.export`	`export_onnx`, `export_torchscript` for deployment outside PyTorch

Examples

The docs/tutorials/ folder has six runnable notebooks, from a basic distillation walkthrough to full deployment:

#	Notebook	Covers
01	Distillation on CIFAR10	`HintonLoss`, `FeatureLoss`, combining losses
02	Distillation on SST-2	Adapting HuggingFace models, `AttentionMapLoss`
03	LLM text generation	`ReverseKLLoss`, mismatched tokenizers, sequence-level KD
04	Compression	`Pruner`, `ChannelPruner`, `Quantizer`
05	Full training control	Custom losses, callbacks, AMP, checkpointing, custom engines
06	Exporting for deployment	`export_onnx`, `export_torchscript`

Contact & Contributing

You can report an issue directly on GitHub. Bug reports, feature requests, and questions are all welcome.

Contributions are welcome too, whether it's a bug fix, a new feature, or a documentation improvement:

Open an issue first for anything non-trivial, to discuss the approach before you start.
Fork the repository and work on a dedicated branch.
Add or update tests for any behavior change (run uv run pytest locally before opening a PR).
Keep the code clean: uv run ruff check and uv run ruff format.
Open a pull request against main, CI runs the test suite, coverage, and lint checks automatically.

Citation

If you use ShrinkAI in your work and think it was helpful, please cite it as:

@software{shrinkai2026,
  author  = {Elouan Marsot},
  title   = {ShrinkAI},
  url     = {https://github.com/elouanzer/shrinkai},
  license = {MIT},
  version = {0.1.0}
}

License

MIT License

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

shrinkai-0.1.0.tar.gz (46.0 kB view details)

Uploaded Sep 13, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

shrinkai-0.1.0-py3-none-any.whl (62.0 kB view details)

Uploaded Sep 13, 2026 Python 3

File details

Details for the file shrinkai-0.1.0.tar.gz.

File metadata

Download URL: shrinkai-0.1.0.tar.gz
Upload date: Sep 13, 2026
Size: 46.0 kB
Tags: Source
Uploaded using Trusted Publishing? Yes
Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for shrinkai-0.1.0.tar.gz
Algorithm	Hash digest
SHA256	`cfcf28320520242819d650949acd62d2d2264d20f6d56c4f0df4878d9c99995a`
MD5	`d24ee7358e1cab0ae8ae80d8df58b70d`
BLAKE2b-256	`e9fde3dbcd55e633e5b4afaf62317c0dfe948e8982ab4dbd655a9fa69c681fb6`

See more details on using hashes here.

Provenance

The following attestation bundles were made for shrinkai-0.1.0.tar.gz:

Publisher: publish.yml on elouanzer/shrinkai

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Statement:
- Statement type: https://in-toto.io/Statement/v1
- Predicate type: https://docs.pypi.org/attestations/publish/v1
- Subject name: shrinkai-0.1.0.tar.gz
- Subject digest: cfcf28320520242819d650949acd62d2d2264d20f6d56c4f0df4878d9c99995a
- Sigstore transparency entry: 2819682730
- Sigstore integration time: Sep 13, 2026
Source repository:
- Permalink: elouanzer/shrinkai@41cedd4ee861703d3ee2caf189f27372b13c3ce4
- Branch / Tag: refs/tags/v0.1.0
- Owner: https://github.com/elouanzer
- Access: public
Publication detail:
- Token Issuer: https://token.actions.githubusercontent.com
- Runner Environment: github-hosted
- Publication workflow: publish.yml@41cedd4ee861703d3ee2caf189f27372b13c3ce4
- Trigger Event: release

File details

Details for the file shrinkai-0.1.0-py3-none-any.whl.

File metadata

Download URL: shrinkai-0.1.0-py3-none-any.whl
Upload date: Sep 13, 2026
Size: 62.0 kB
Tags: Python 3
Uploaded using Trusted Publishing? Yes
Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for shrinkai-0.1.0-py3-none-any.whl
Algorithm	Hash digest
SHA256	`acab908ff816f4bbe1f753bfab05ff3b7b84b0a8a7e4fae20a29407f241edb14`
MD5	`37836a046b5107ed238836ad8ab24957`
BLAKE2b-256	`e7832cb1c077d3d91aae731f1f376010dfb2d504ea3ad27edb058c7602899949`

See more details on using hashes here.

Provenance

The following attestation bundles were made for shrinkai-0.1.0-py3-none-any.whl:

Publisher: publish.yml on elouanzer/shrinkai

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Statement:
- Statement type: https://in-toto.io/Statement/v1
- Predicate type: https://docs.pypi.org/attestations/publish/v1
- Subject name: shrinkai-0.1.0-py3-none-any.whl
- Subject digest: acab908ff816f4bbe1f753bfab05ff3b7b84b0a8a7e4fae20a29407f241edb14
- Sigstore transparency entry: 2819682805
- Sigstore integration time: Sep 13, 2026
Source repository:
- Permalink: elouanzer/shrinkai@41cedd4ee861703d3ee2caf189f27372b13c3ce4
- Branch / Tag: refs/tags/v0.1.0
- Owner: https://github.com/elouanzer
- Access: public
Publication detail:
- Token Issuer: https://token.actions.githubusercontent.com
- Runner Environment: github-hosted
- Publication workflow: publish.yml@41cedd4ee861703d3ee2caf189f27372b13c3ce4
- Trigger Event: release

Release history Release notifications | RSS feed

This release

0.1.0 This release

Sep 13, 2026 2 files

shrinkai 0.1.0

ShrinkAI

Installation

Quick Start

Architecture

Examples

Contact & Contributing

Citation

License

Project links

Key dates

1 maintainer

Credits

GitHub Statistics

License expression

Requires

Provides Extra

Tags

Classifiers

Download files

Source Distribution

Built Distribution

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

Release history Release notifications | RSS feed