By Yikai Wang, Yi Yang, Fuchun Sun, Anbang Yao.

This pytorch implementation is a two-stage improved version of our paper "Sub-bit Neural Networks: Learning to Compress and Accelerate Binary Neural Networks", published in ICCV 2021.

Following this repository,

• Download the ImageNet dataset from https://www.image-net.org.
• Then move validation images to labeled subfolders, using the following script.

• python3, pytorch 1.4.0, torchvision 0.5.0

(1) Step1: binarizing activations (or you can omit this step by using our Step1 model checkpoint_ba.pth.tar),

• Change directory to ./step1,
• Run the folowing script,

CUDA_VISIBLE_DEVICES=0,1,2,3 python train.py --data=path/to/ILSVRC2012/  --batch_size=512 --learning_rate=1e-3 --epochs=256 --weight_decay=1e-5

(2) Step2: binarizing weights + activations,

• Change directory to ./step2,
• Create new folder ./models and copy checkpoint_ba.pth.tar (obtained from Step1) to ./models,
• Run the folowing script,

CUDA_VISIBLE_DEVICES=0,1,2,3 python train.py --data=path/to/ILSVRC2012/  --batch_size=512 --learning_rate=1e-3 --epochs=256 --weight_decay=0 --bit-num=5

• Comment: --bit-num=5 corresponds to 0.56 bit (bit-num indicates tau in the paper).

This implementation is based on ResNet-18 of ReActNet.

If you find our code useful for your research, please consider citing:

@inproceedings{wang2021snn,
    title={Sub-bit Neural Networks: Learning to Compress and Accelerate Binary Neural Networks},
    author={Wang, Yikai and Yang, Yi and Sun, Fuchun and Yao, Anbang},
    booktitle={International Conference on Computer Vision (ICCV)},
    year={2021}
}

SNN is released under MIT License.