This repository contains the code to reproduce experiments in the paper "Phase Collapse in Neural Networks", accepted at ICLR 2022.

Our code is designed to run on GPU using PyTorch. In order to run our experiments, you will need the following packages: numpy, scipy, torch (1.8), and torchvision. See requirements.txt for the precise (but not minimal) environment which was used to run the experiments in the paper, it can be installed with pip install -r requirements.txt.

The ImageNet dataset must be downloaded from https://www.image-net.org or from https://www.kaggle.com/c/imagenet-object-localization-challenge/overview/description ((registration required). Then move validation images to labeled subfolders, using the PyTorch shell script.

CIFAR-10 is automatically downloaded by Pytorch.

To train a model, run main_block.py with the desired arguments. Running run.py will train all models whose results are reported in Tables 1, 2 and 3, if provided with the path to the ImageNet dataset.