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}
}