Skip to content

Latest commit

 

History

History
82 lines (63 loc) · 3.83 KB

README.md

File metadata and controls

82 lines (63 loc) · 3.83 KB
Raw

Code for Normalizer-Free Networks

This repository contains code for the ICLR 2021 paper "Characterizing signal propagation to close the performance gap in unnormalized ResNets," by Andrew Brock, Soham De, and Samuel L. Smith, and the arXiv preprint "High-Performance Large-Scale Image Recognition Without Normalization" by Andrew Brock, Soham De, Samuel L. Smith, and Karen Simonyan.

Running this code

Using run.sh will create and activate a virtualenv, install all necessary dependencies and run a test program to ensure that you can import all the modules and take a single experiment step. To train with this code, use this virtualenv and use one of the experiment.py files in combination with JAXline. The provided demo Colab can be run online, or by starting a jupyter notebook within this virtualenv.

Note that you will need a local copy of ImageNet compatible with the TFDS format used in dataset.py in order to train on ImageNet.

Pre-Trained Weights

We provide pre-trained weights for NFNet-F0 through F5 (trained without SAM), and for NFNet-F6 trained with SAM. All models are pre-trained on ImageNet for 360 epochs at batch size 4096, and are provided as numpy files containing parameter trees compatible with haiku. In utils.py we provide a load_haiku_file function which loads these parameter trees, and flatten_haiku_tree to convert these to flat dictionaries which may prove easier to port to other frameworks. Note that we do not provide model states, as these models, lacking batchnorm, do not have running stats. Note also that the conv layer weights are in the format HWIO, so for frameworks like PyTorch which use OIHW you'll need to swap the axes appropriately to the layout you use.

Model #FLOPs #Params Top-1 Top-5 TPUv3 Train GPU Train link
F0 12.38B 71.5M 83.6 96.8 73.3ms 56.7ms haiku
F1 35.54B 132.6M 84.7 97.1 158.5ms 133.9ms haiku
F2 62.59B 193.8M 85.1 97.3 295.8ms 226.3ms haiku
F3 114.76B 254.9M 85.7 97.5 532.2ms 524.5ms haiku
F4 215.24B 316.1M 85.9 97.6 1033.3ms 1190.6ms haiku
F5 289.76B 377.2M 86.0 97.6 1398.5ms 2177.1ms haiku
F6+SAM 377.28B 438.4M 86.5 97.9 2774.1ms - haiku

Demo Colab Open In Colab

We also include a Colab notebook with a demo showing how to run an NFNet to classify an image.

Giving Credit

If you use this code in your work, we ask you to please cite one or both of the following papers.

The reference for the Normalizer-Free structure and NF-ResNets or NF-Regnets:

@inproceedings{brock2021characterizing,
  author={Andrew Brock and Soham De and Samuel L. Smith},
  title={Characterizing signal propagation to close the performance gap in
  unnormalized ResNets},
  booktitle={9th International Conference on Learning Representations, {ICLR}},
  year={2021}
}

The reference for Adaptive Gradient Clipping (AGC) and the NFNets models:

@article{brock2021high,
  author={Andrew Brock and Soham De and Samuel L. Smith and Karen Simonyan},
  title={High-Performance Large-Scale Image Recognition Without Normalization},
  journal={arXiv preprint arXiv:2102.06171},
  year={2021}
}