NeuralZip is a out-of-box Python scaffold for neural network quantization-aware training (QAT). With NeuralZip, you can focus on optimizing your quantization algorithm itself without falling into error-prone and dirty jobs.

NeuralZip can:

• Inject quantizer operators into your network without modification
• Decouple your quantizer implementation with the network implementation
• Built-in quantizers, including LSQ/...
• Evaluate your quantization algorithm on common benchmarks, including ImageNet/CIFAR10/...
• Automatic accelerate training with multi-thread DistributedDataParallel
• Checkpoint / TensorBoard visualization / YAML & CLI configuration / Thorough logs

Here are some LSQ-based experiment results. Their configuration files can be found in the conf folder.

Please view this link for more details (such as experiment setup and CLI commands).

Install library dependencies within an Anaconda environment.

# Create a environment with Python 3.10
conda create -n neuralzip python=3.10
conda activate neuralzip
# PyTorch GPU version 1.13
conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
# PyTorch Lightning & its Bolts
python -m pip install --upgrade lightning lightning-bolts torchmetrics
# Miscellaneous
conda install omegaconf

This program use YAML files as inputs. A template as well as the default configuration is provided as conf/template.yaml. Please read it before running Python scripts. It is disallowed to modify this file and use it as your configuration, which may cause severe errors at the runtime.

If you want to change the behaviour of this program, please copy it somewhere else. And then run the main.py with your modified configuration file.

python main.py conf_filepath=/path/to/your/config/file.yaml

The modified options in your YAML file will overwrite the default settings. For details, please read the comments in conf/template.yaml. You can also find some example configuration files in the example folder.

You can also use CLI arguments to configure the program, like this:

python main.py conf_filepath=some.yaml optimizer.lr=0.05

Thus, the setting optimizer.lr in some.yaml will be overridden with the CLI input (i.e. 0.05).

If you find any bugs in my code or have any ideas to improve the quantization results, please feel free to open an issue. I will be glad to join the discussion.

NeuralZip originates from my another project, an implementation of LSQ-Net.

It is powered by PyTorch, PyTorch-Lightning and many other open-source projects. Thanks for their excellent jobs.