This repository contains the source code for a specialized framework focused on mixed-precision quantization of neural networks. The framework leverages NSGA-II for searching optimal quantization configurations and an extended version of Timeloop to model mixed-precision data types for hardware metrics retrieval.
README.md
- README description of the frameworkrequirements.yml
- requirements file for setting up the necessary environmentrun_nsga.sh
- shell script for running the nsga experimentsrc/
- the main directory containing all the source coderun_nsga.py
- script interface to initiate and run the nsga processmapper_facade.py
- interface script for calling timeloop-mapper and caching hardware resultsnsga/
- contains scripts related to NSGA-II algorithmpytorch/
- includes all scripts related to PyTorch models, data management, training, evaluation, and utilities implementationtimeloop_utils/
- contains Timeloop interface configurations and helper scripts
$ conda env create --file requirements.yml
$ conda activate qatizer
$ pip install -r requirements.yml
- IMPORTANT NOTE:
mapper_facade.py
requires Timeloop+Accelergy infrastructure to be installed. If it is not installed, you will need to use docker or load cached HW metrics – in that case, modifications tomapper_facade.py
might be necessary.
- Please download your desired datasets into the
src/pytorch/data/datasets/
folder. See README instructions there for file structure information.
run_nsga.sh
- Script to run NSGA-II algorithm, tailored for mixed-precision quantization tasks (MODIFY AS YOU WISH).
src/pytorch/run_scripts/
- For (quantization-aware) training and or evaluation of a PyTorch model, you can create scripts to run thetrain.py
andrun.py
scripts located insrc/pytorch/
.