STTN

Paper: Soft Threshold Ternary Networks

Install Ternary Activation CUDA Extension

• This repo uses Pytorch 1.7.1+cu101
• Our CUDA extension are build with gcc version 5.3.1

cd TerActivateFunc
sh install.sh

Training Time: Train STTN (ternarized W and A) on ImageNet

python main.py --data path_to_ImageNet --gpu-id 0,1,2,3 -j16 --arch resnet18_bwx2 --save=res18_bwx2_ta

After training, two parallel binary weights are obtained.

Inference Time: Merge and Evaluate

Previous ternary networks introduces a hard threshold to divide quantization intervals, which sets an additional constraint and limits the ability of ternary networks:

STTN enables the model to automatically determine which weights to be -1/0/1, avoiding the hard threshold . Concretely, at training time, we replace the ternary convolution filter with two parallel binary convolution filters and . They are both binary-valued and have the same shape with ternary filter. Due to the additivity of convolutions with the same kernel sizes, a new kernel can be obtained by: . During inference, the ternary weights are used.

• We provide a simple script to merge the two parallel binary weights into an equivalent ternary weight.

python merge_from_bwx2_to_ternary.py

• You can test the accuracy with the merged ternary weights. It should be the same as the result of two parallel binary weights during training, because they are equivalent.

python evaluate.py --data path_to_ImageNet -j16 --arch resnet18_ter --resume ./ternary/resnet18_ter.pth.tar --evaluate

Results on ResNet-18

Model	Weight	Activation	Accuracy	Download
Float	FP32	FP32	69.3
STTN (paper report)	ternary	ternary	66.2%
STTN (this repo)	ternary	ternary	68.2%	google cloud

Citation

@inproceedings{xu2021soft,
  title={Soft threshold ternary networks},
  author={Xu, Weixiang and He, Xiangyu and Zhao, Tianli and Hu, Qinghao and Wang, Peisong and Cheng, Jian},
  booktitle={Proceedings of the Twenty-Ninth International Conference on International Joint Conferences on Artificial Intelligence},
  pages={2298--2304},
  year={2021}
}

@ARTICLE{9927185,
  author={Xu, Weixiang and Li, Fanrong and Jiang, Yingying and Yong, A and He, Xiangyu and Wang, Peisong and Cheng, Jian},
  journal={IEEE Transactions on Circuits and Systems for Video Technology}, 
  title={Improving Extreme Low-bit Quantization with Soft Threshold}, 
  year={2022},
  volume={},
  number={},
  pages={1-1},
  doi={10.1109/TCSVT.2022.3216389}}