The notebook neuralcollapse/neuralcollapse.ipynb
contains code demonstrating Neural Collapse on Cross-Entropy and MSE Loss.
Notebook is designed to be short, easy-to-interpret, and executable
from the browser using Google Colab.
(There is a link in the notebook to open in Colab.)
MNIST-ResNet18 was chosen because it ran most reliably within the in-browser the memory constraints of Google Colab. If you are still getting out-of-memory errors, try clicking "Runtime"->"Factory Reset Runtime" on the Colab menu bar.
It should be clear how to adapt the code to other networks-dataset combinations to be run on local clusters with more memory. Happy experimenting!