# PrivCirNet: Efficient Private Inference via Block Circulant Transformation

<div align="center">

</div>

#### 🔥 Introducing a protocol/network co-optimization framework for efficient homomorphic encryption-based private DNN inference
#### 🔥 PrivCirNet supports representative DNN architecture including MobileNetV2, ResNet-18 and Vision Transformer and datasets (CIFAR-10/100, Tiny ImageNet, ImageNet)
## PrivCirNet zoo
We provide PrivCirNet models which can be downloaded here **(weights are in the `pretrained` folder, we only provide typical checkpoints due to size limitations of supplementary materials)**. Once the paper is accepted, we will release the Google Drive link to all checkpoints.

## Installation

1. Install `torch>=2.0.0`, `python>=3.10`.
2. Install other pip packages via `pip install -r requirements.txt`.
3. Prepare the dataset including CIFAR-10, CIFAR100, Tiny ImageNet and ImageNet.

## Fast test PrivCirNet's models
Download checkpoints and put them in the `pretained` folder. Change the `initial_checkpoint` in the `.yml` file to the path of the checkpoint. Then run the same command as training scripts `train_cir.sh` to test the models. For example, to test ViT on Tiny ImageNet, you can run the following command:
```shell
CUDA_VISIBLE_DEVICES=0 python train_cir.py -c configs/datasets/ViT/tiny_fix.yml --model vit_9_12_64 PATH_TO_TINY_IMAGENET
```
## Training Scripts
### Step1: Run CirILP.py to get block size configuration
This step is simple and the script is `ILP.sh`.\
For example, to configure ViT on CIFAR-10, you can run the following command:
```shell
CUDA_VISIBLE_DEVICES=1 python CirILP.py -c configs/datasets/ViT/cifar10_ILP.yml --model vit_7_4_32 PATH_TO_CIFAR10
```
A log file named `vit_c10_ILP.log` will be created to save the configuration and logs. You will get the layer-wise block sizes configuration in the log file.

There are two parameters in the `.yml` file you need to specify:
- budget. You can set the budget to `2`, `4` and `8` whose latency is less than that of networks with uniform block size `2`, `4` and `8`.
- better_initialization. `True` means use the initialization method proposed in PrivCirNet, `False` means use the previous initialization method, i.e., $\min |W'-W|_2^2$.

### Step2: Train the circulant models
This step is also simple which is the same as training the original models. The script is `train_cir.sh`. For example, to train ViT on CIFAR-10, you can run the following command:
```shell
CUDA_VISIBLE_DEVICES=1 python train_cir.py -c configs/datasets/ViT/cifar10_fix.yml --model vit_7_4_32 PATH_TO_CIFAR10
```
A log file named `your_log_name.log` will be created to save logs. The checkpoints will be saved in the `output` folder. There are several parameters in the `.yml` file you need to specify:
- `fix_blocksize_list`. It is the block size configuration you get from the previous step. Must be separated by commas.
- `log_name`. The name of the log file.
- `use_kd`. Whether to use knowledge distillation. If `True`, you need to specify the `teacher_model` and `teacher_checkpoint` in the `.yml` file.
- `initial_checkpoint`. You must load the initial checkpoint from the original model. All initial checkpoints can be downloaded here (weights are in the `pretrained` folder).

### Train other baseline models
This repository also provides the training scripts for the two baselines:
- Uniform block size circulant networks. The training script is the same as `Train the circulant models` where you can set `fix_blocksize` to `2`,`4` and `8` to train uniform block size circulant networks.
- SpENCNN (structured pruning). The training script is `train_prune.sh`. For example, to train MobileNetV2 on CIFAR-10, you can run the following command:
```shell
CUDA_VISIBLE_DEVICES=1 python train_prune.py -c configs/datasets/Prune/cifar10.yml --model c10_prune_mobilenetv2 PATH_TO_CIFAR10
```
- You can set `prune_ratio` in the `.yml` file to specify the pruning ratio.
- Pruning models can be downloaded here (weights are in the `pretrained` folder).

### Train the original models
This repository also provides the training scripts for the original models although we provide the pretrained checkpoints. The training script is `train.sh`. For example, to train ResNet-18 on Tiny ImageNet, you can run the following command:
```shell
CUDA_VISIBLE_DEVICES=0 python train.py -c configs/datasets/ResNet/tiny.yml --model tiny_resnet18 PATH_TO_TINY_IMAGENET
```
## License
This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.


## Citation
If our work assists your research, feel free to give us a star ⭐ or cite us using:
```
Coming soon
```