Home

Awesome

SPIN

This repository contains the official implementation for the ECCV'22 paper, "SPIN: An Empirical Evaluation on Sharing Parameters of Isotropic Networks".

Code Overview

We provide the implementation of weight sharing version of the ConvMixer model. The main code for the implementation are in the models directory. The model can be configured by the files in configs. We provide three example configs.

Note that in order to run the model configs/WF-WSConvMixer.yaml, you must have a corresponding pretrained ConvMixer model. Please refer to our paper for each technique.

Installation

First, clone this repo with

git clone https://github.com/apple/ml-spin.git

The implementation of SPIN reuses the infrastructure of Meta Research's open source project SlowFast. Our modification to the SlowFast code is stored in the spin-slowfast.patch. To download the SlowFast code and apply our changes, run

bash setup.sh

After getting the codebase ready, follow this link from SlowFast repo to setup your environment and install other dependencies.

Training

After the environment is set up, you can run the following example training script to train a weight sharing ConvMixer model. The script assumes you have a machine with 4-GPUs.

bash run.sh

Pre-trained ConvMixer Models on ImageNet1K

We provide our pretrained models of ConvMixer, WS-ConvMixer and WFWS-ConvMixer in the following table. For the WFWS-ConvMixer, we first initialized the model using the proposed weight fusion technique with mean strategy, and then run the models/fuse_weights.py to export the fused model after training. In order to re-run the model, please use the WS-ConvMixer configuration. Please note we did a light hyperparameter tunning so the accuracy is slightly higher than the numbers reported in the paper.

C/D/P/KWeight Sharing?Weight Fusion?Sharing RateShare DistributionSharing MappingAccuracyModel Size
768/32/14/3NoNo---76.32%79MB
768/32/14/3YesNo2UniformSequential74.27%43MB
768/32/14/3YesMean2UniformSequential75.21%43MB

Citation

If you find our code or paper helps, please consider citing:

@article{spin_eccv22,
    author    = {Lin, Chien-Yu and Prabhu, Anish and Merth, Thomas and Mehta, Sachin and Ranjan, Anurag and Horton, Maxwell and Rastegari, Mohammad}
    title     = {SPIN: An Empirical Evaluation on Sharing Parameters of Isotropic Networks},
    booktitle = {Proceedings of the European Conference on Computer Vision (ECCV)},
    year      = {2022}
}