mkotyushev/otfusion

Model Fusion via Optimal Transport

Requirements

Install the Python Optimal Transport Library

pip install POT

Other than that, we also need PyTorch v1 or higher and NumPy. (Also, Python 3.6 +)

Before running, unzip the respective pretrained model zip file. Also, you need to unzip the cifar.zip file for some imports to work.

Sample commands of one-shot model fusion

For MNIST + MLPNet

python main.py --gpu-id 1 --model-name mlpnet --n-epochs 10 --save-result-file sample.csv \
--sweep-name exp_sample --exact --correction --ground-metric euclidean --weight-stats \
--activation-histograms --activation-mode raw --geom-ensemble-type acts --sweep-id 21 \
--act-num-samples 200 --ground-metric-normalize none --activation-seed 21 \
--prelu-acts --recheck-acc --load-models ./mnist_models --ckpt-type final \
--past-correction --not-squared --dist-normalize --print-distances --to-download

For CIFAR10 + VGG11

python main.py --gpu-id 1 --model-name vgg11_nobias --n-epochs 300 --save-result-file sample.csv \
--sweep-name exp_sample --correction --ground-metric euclidean --weight-stats \
--geom-ensemble-type wts --ground-metric-normalize none --sweep-id 90 --load-models ./cifar_models/ \
--ckpt-type best --dataset Cifar10 --ground-metric-eff --recheck-cifar --activation-seed 21 \
--prelu-acts --past-correction --not-squared --normalize-wts --exact

We also recommend that users play around with some of options or hyper-parameters above, as the commands listed here are not highly tuned. For instance, getting rid of the --normalize-wts flag and running the below command instead, results in a test accuracy of 86.51% instead of 85.98% on CIFAR10.

python main.py --gpu-id 1 --model-name vgg11_nobias --n-epochs 300 --save-result-file sample.csv \
--sweep-name exp_sample --correction --ground-metric euclidean --weight-stats \
--geom-ensemble-type wts --ground-metric-normalize none --sweep-id 90 --load-models ./cifar_models/ \
--ckpt-type best --dataset Cifar10 --ground-metric-eff --recheck-cifar --activation-seed 21 \
--prelu-acts --past-correction --not-squared --exact

For CIFAR10 + ResNet18

python main.py --gpu-id 1 --model-name resnet18_nobias_nobn --n-epochs 300 --save-result-file sample.csv \
--sweep-name exp_sample --exact --correction --ground-metric euclidean --weight-stats \
--activation-histograms --activation-mode raw --geom-ensemble-type acts --sweep-id 21 \
--act-num-samples 200 --ground-metric-normalize none --activation-seed 21 --prelu-acts --recheck-acc \
--load-models ./resnet_models/ --ckpt-type best --past-correction --not-squared  --dataset Cifar10 \
--handle-skips

The code and pretrained models correspond to the paper: Model Fusion via Optimal Transport. If you use any of the code or pretrained models for your research, please consider citing the paper as.

@article{singh2020model,
  title={Model fusion via optimal transport},
  author={Singh, Sidak Pal and Jaggi, Martin},
  journal={Advances in Neural Information Processing Systems},
  volume={33},
  year={2020}
}