S2AP: Score-space Sharpness Minimization for Adversarial Pruning

This repository contains the official implementation of S2AP (Score-Sharpness aware Adversarial Pruning), a method designed to minimize sharpness, improve adversarial robustness, and stabilize mask search during network pruning.

The code supports training, pruning, sharpness evaluation, and adversarial robustness evaluation.

---

📦 Setup

Install required packages with:

pip install -r requirements.txt

---

Step 1: Pretrain a dense model

Run pretraining

python train.py --arch resnet18 --exp-mode pretrain --exp-name resnet18 --layer-type subnet --k 1.0 --dataset CIFAR10 --data-dir /path/to/shared_data --batch-size 128 --optimizer sgd --trainer adv --adv_loss trades --val_method adv --gpu 0 --epochs 100

It is possible to select different architectures, such as vgg16_bn and wrn_28_4 Similarly, one can choose between CIFAR10, SVHN, or ImageNet (for ImageNet, it is recommendable to choose fat as a trainer) as datasets, as well as using different losses among the ones available.

Evaluate pretrained

python aa_eval.py --model-path /path/to/resnet18/pretrain/model_best_dense.pth.tar --dataset CIFAR10 --model resnet18_dense --gpu 0

---

Step 2: Prune the pretrained models and evaluate mask sharpness and robustness

Prune with Orig. method

python train.py --arch resnet18 --source-net /path/to/resnet18/pretrain/pretrained.pth.tar --exp-mode harp_prune --exp-name resnet18/adv/80 --layer-type subnet --k 0.1 --prune_reg 'weight' --stg_id '010_uni' --dataset CIFAR10 --data-dir /path/to/shared_data --batch-size 128 --optimizer sgd --trainer adv --adv_loss trades --val_method adv --gpu 0 --epochs 20 --hamming --eigenvalues

Prune with S2AP

python train.py --arch resnet18 --source-net /path/to/resnet18/pretrain/model_best.pth.tar --exp-mode harp_prune --exp-name resnet18/s2ap/80 --layer-type subnet --k 0.1 --prune_reg 'weight' --stg_id '010_uni' --dataset CIFAR10 --data-dir /path/to/shared_data --batch-size 128 --optimizer sgd --trainer s2ap --s2ap-warmup 5 --s2ap-gamma 0.001 --adv_loss trades --val_method adv --gpu 0 --epochs 20 --hamming --eigenvalues

In this step, it is likewise possible to prune with HYDRA by setting --exp-mode score_prune. To control sparsity instead, one can use the parameter k (e.g., --k 0.1 for sparsity rate 90%). The arguments --hamming and --eigenvalues instead, allow saving the different pruning masks to compute hamming distance and saving the eigenvalues to later compute the maximum eigenvalue. Overall, using S2AP amounts to using the s2ap trainer and specifying warm-up epochs and perturbation through, respectively --s2ap-warmup and --s2ap-gamma. When using ImageNet, it is instead important to specify s2ap_fat as a trainer, while the rest can be tuned accordingly.

Eval pruned model robustness (Orig.)

python aa_eval.py --model-path /path/resnet18/adv/80/model_best_dense.pth.tar --dataset CIFAR10 --model resnet18_dense --gpu 0

Eval pruned model robustness (S2AP)

python aa_eval.py --model-path /path/resnet18/s2ap/80/model_best_dense.pth.tar --dataset CIFAR10 --model resnet18_dense --gpu 0

Eval pruned model sharpness (Orig.)

python sharpness/eval_sharpness.py --k 0.2 --task_mode harp_finetune --dataset=cifar10 --model=resnet18 --n_eval_sharpness=1024 --bs_sharpness=128 --adaptive --normalize_logits --n_iters=20 --algorithm=m_apgd_linf --rho=0.008 --log_folder='my_folder' --sparse --masked --model_path="/path/to/resnet18/adv/80/model_best_dense.pth.tar" --gpu 0 --sharpness_on_test_set;

Eval pruned model sharpness (S2AP)

python sharpness/eval_sharpness.py --k 0.2 --task_mode harp_finetune --dataset=cifar10 --model=resnet18 --n_eval_sharpness=1024 --bs_sharpness=128 --adaptive --normalize_logits --n_iters=20 --algorithm=m_apgd_linf --rho=0.008 --log_folder='my_folder' --sparse --masked --model_path="/path/to/resnet18/s2ap/80/model_best_dense.pth.tar" --gpu 0 --sharpness_on_test_set;

This section finally allows evaluating the adversarial robustness of the masks and computing their sharpness. While doing so, one can set the perturbation added to perturb the score-space and measure sharpness through the --rho parameter.

---

Step 3: Finetune pruned models and evaluate robustness

Finetune pruned model (Orig.)

python train.py --arch resnet18 --source-net /path/to/resnet18/adv/80/model_best.pth.tar --exp-mode harp_finetune --exp-name resnet18/adv/80 --layer-type subnet --k 0.1 --dataset CIFAR10 --data-dir /.../shared_data --batch-size 128 --optimizer sgd --trainer adv --adv_loss trades --val_method adv --gpu 0 --epochs 100

Finetune pruned model (S2AP)

python train.py --arch resnet18 --source-net /path/to/resnet18/s2ap/80/model_best.pth.tar --exp-mode harp_finetune --exp-name resnet18/s2ap/80 --layer-type subnet --k 0.1 --dataset CIFAR10 --data-dir /.../shared_data --batch-size 128 --optimizer sgd --trainer s2ap --s2ap-warmup 10 --s2ap-gamma 0.001 --adv_loss trades --val_method adv --gpu 0 --epochs 100

Evaluate finetuned model (Orig.)

python aa_eval.py --model-path /path/to/resnet18/adv/80/model_best_dense.pth.tar --dataset CIFAR10 --model resnet18_dense --gpu 0

Evaluate finetuned model (S2AP)

python aa_eval.py --model-path /path/to/resnet18/s2ap/80/model_best_dense.pth.tar --dataset CIFAR10 --model resnet18_dense --gpu 0

---

Compute Mask Stability and Largest Eigenvalues

For this step, you can run the files plots/mask_stability.py and plots/eigenvalues.py, upon selecting the appropriate datasets, models, sparsities, and pruning methods. This will allow reproducing the plots presented in the paper.