CNN Wild Park - Graph Neural Networks for Learning Equivariant Representations of Neural Networks

This repository contains the CNN Wild Park dataset from the paper:

Graph Neural Networks for Learning Equivariant Representations of Neural Networks
Miltiadis Kofinas*, Boris Knyazev, Yan Zhang, Yunlu Chen, Gertjan J. Burghouts, Efstratios Gavves, Cees G. M. Snoek, David W. Zhang*
ICLR 2024 (oral)
https://arxiv.org/abs/2403.12143
https://github.com/mkofinas/neural-graphs
*Joint first and last authors

We introduce a new dataset of CNNs, which we term CNN Wild Park.
The dataset consists of 117,241 checkpoints from 2,800 CNNs, trained for up to 1,000 epochs on CIFAR10.
The CNNs vary in the number of layers, kernel sizes, activation functions, and residual connections between arbitrary layers.

More specifically, we construct the CNN Wild Park dataset by training 2,800 small CNNs with different architectures for 200 to 1,000 epochs on CIFAR10. We retain a checkpoint of its parameters every 10 steps and also record the test accuracy. The CNNs vary by:

• Number of layers L in [2, 3, 4, 5] (note that this does not count the input layer).
• Number of channels per layer c_l in [4, 8, 16, 32].
• Kernel size of each convolution k_l in [3, 5, 7].
• Activation functions at each layer are one of ReLU, GeLU, tanh, sigmoid, leaky ReLU, or the identity function.
• Skip connections between two layers with at least one layer in between. Each layer can have at most one incoming skip connection. We allow for skip connections even in the case when the number of channels differ, to increase the variety of architectures and ensure independence between different architectural choices. We enable this by adding the skip connection only to the min(c_n, c_m) nodes.

We divide the dataset into train/val/test splits such that checkpoints from the same run are not contained in both the train and test splits.