Skip to content

BitPack is a practical tool to efficiently save ultra-low precision/mixed-precision quantized models.

License

Notifications You must be signed in to change notification settings

Zhen-Dong/BitPack

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

19 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

BitPack

BitPack is a practical tool that can efficiently save quantized neural network models with mixed bitwidth.

Installation

  • PyTorch version >= 1.4.0
  • Python version >= 3.5
  • To install Bitpack simply run:
git clone https://github.com/Zhen-Dong/BitPack.git
cd BitPack

Usage

  • We can use BitPack pack.py to save integer checkpoints with various bitwidth, and use BitPack unpack.py to load the packed checkpoint, as shown in the demo.
  • To pack integer values that are saved in floating point format, add --force-pack-fp in the command.
  • To directly save packed checkpoint in PyTorch, please use save_quantized_state_dict() and load_quantized_state_dict() in pytorch_interface.py. If you don't want to operate jointly on state_dict, then codes inside the for loop of those two functions can be applied on every quantized tensor (ultra low-precision integer tensors) in various quantization frameworks.



Quick Start

BitPack is handy to use on various quantization frameworks. Here we show a demo that applying BitPack to save mixed-precision model generated by HAWQ.

export CUDA_VISIBLE_DEVICES=0
python pack.py --input-int-file quantized_checkpoint.pth.tar --force-pack-fp
python unpack.py --input-packed-file packed_quantized_checkpoint.pth.tar --original-int-file quantized_checkpoint.pth.tar

To get a better sense of how BitPack works, we provide a simple test that compares the original tensor, the packed tensor, and the unpacked tensor in details.

cd bitpack
python bitpack_utils.py

Results of BitPack on ResNet50

Original Precision Quantization Original Size(MB) Packed Size(MB) Compression Ratio
Floating Point Mixed-Precision(4bit/8bit) 102 13.8 7.4x
8-bit Mixed-Precision(2bit/8bit) 26 7.9 3.3x

Special Notes

  • unpack.py can be used for checking correctness. It loads and unpacks the packed model, and then compares it with the original model.

License

THIS SOFTWARE WAS DEPOSITED IN THE BAIR OPEN RESEARCH COMMONS REPOSITORY ON FEB 1, 2023.

BitPack is released under the MIT license.

About

BitPack is a practical tool to efficiently save ultra-low precision/mixed-precision quantized models.

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Contributors 4

  •  
  •  
  •  
  •  

Languages