GradMax: Growing Neural Networks using Gradient Information

Code for reproducing our results in the GradMax paper [arxiv.org/abs/2201.05125].

Setup

First clone this repo.

git clone https://github.com/google-research/growneuron.git
cd growneuron

Following script installs the necessary libraries and runs few tests.

bash run.sh

Following will download the data and run the baseline experiment. If the data is already downloaded use the --data_dir flag to pass the path.

python growneuron/cifar/main.py --output_dir=/tmp/cifar --download_data

Running GradMax

Following command would start a training with WRN-28-0.25x and grow it into WRN-28-1. Growth is done every 2500 step starting from iteration 10000 at all convolutional layers at once.

rm -rf /tmp/cifar
python growneuron/cifar/main.py --output_dir=/tmp/cifar \
--config=growneuron/cifar/configs/grow_all_at_once.py \
--config.grow_type=add_gradmax

Other Experiments

• Baselines for WRN-28 and VGG11 can be ran using the corresponding configs in growneuron/cifar/configs/.
• Set --config.grow_type argument to add_gradmax, add_firefly, add_gradmax_opt or add_random to grow using different strategies.
• Use --config.is_outgoing_zero to run experiments where outgoing weights are set to zero.
• Use --config.model.normalization_type=batchnorm to run experiments with batch normalization layers.

Disclaimer

This is not an officially supported Google product.