Updated repository

Author: iurada

.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+env/
+record/
+wandb/
+.DS_Store

README.md

@@ -1,2 +1,60 @@
 # Finding Lottery Tickets in Vision Models via Data-driven Spectral Foresight Pruning [CVPR 2024]
 Official code of our work "Finding Lottery Tickets in Vision Models via Data-driven Spectral Foresight Pruning" accepted at CVPR 2024.
+
+<p align="center"><img width="50%" src="./assets/teaser.png"></p>
+
+## Introduction
+<i>Recent advances in neural network pruning have shown how it is possible to reduce the computational costs and memory demands of deep learning models before training. We focus on this framework and propose a new pruning at initialization algorithm that leverages the Neural Tangent Kernel (NTK) theory to align the training dynamics of the sparse network with that of the dense one. Specifically, we show how the usually neglected data-dependent component in the NTK's spectrum can be taken into account by providing an analytical upper bound to the NTK's trace obtained by decomposing neural networks into individual paths. This leads to our Path eXclusion (PX), a foresight pruning method designed to preserve the parameters that mostly influence the NTK's trace. PX is able to find lottery tickets (i.e. good paths) even at high sparsity levels and largely reduces the need for additional training. When applied to pre-trained models it extracts subnetworks directly usable for several downstream tasks, resulting in performance comparable to those of the dense counterpart but with substantial cost and computational savings.</i>
+
+# Setting up the environment
+### Requirements
+Make sure to have a CUDA capable device, running at learst CUDA 11.7. Throughout our experiments we used Python version 3.10.9
+
+### General Dependencies
+To install all the required dependencies go to the root folder of this project and run:
+```bash
+pip install -r requirements.txt
+```
+
+### Datasets
+1. CIFAR-10, CIFAR-100 and Tiny-ImageNet datasets will be downloaded automatically to the folder specified by the `CONFIG.dataset_args['root']` argument once you run the experiments.
+2. ImageNet needs to be downloaded from the official website `https://www.image-net.org/`.
+3. Pascal VOC2012 can be downloaded by running the following script:
+```python
+import os
+import tarfile
+from torchvision.datasets.utils import download_url
+
+db = {'2012': {
+        'url': 'http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar',
+        'filename': 'VOCtrainval_11-May-2012.tar',
+        'md5': '6cd6e144f989b92b3379bac3b3de84fd',
+        'base_dir': 'VOCdevkit/VOC2012'
+    }}
+
+def download_extract(url, root, filename, md5):
+    download_url(url, root, filename, md5)
+    with tarfile.open(os.path.join(root, filename), "r") as tar:
+        tar.extractall(path=root)
+
+if __name__ == '__main__':
+    download_extract(db['2012']['url'], 'data/VOC2012', db['2012']['filename'], db['2012']['md5'])
+```
+
+At this point you should be able to run the provided code.
+
+## Running The Experiments
+Please refer to the `parse_args.py` file for the full list of command line arguments available. You can find in the `launch_scripts/` folder some example scripts used to run the experiments.
+
+## Acknowledgement
+Our code is developed starting from the [Synflow](https://arxiv.org/abs/2006.05467) code repository: https://github.com/ganguli-lab/Synaptic-Flow.
+
+# Citation
+```
+@inproceedings{iurada2024finding,
+  author={Iurada, Leonardo and Ciccone, Marco and Tommasi, Tatiana},
+  booktitle={CVPR}, 
+  title={Finding Lottery Tickets in Vision Models via Data-driven Spectral Foresight Pruning}, 
+  year={2024}
+}
+```

assets/teaser.png

@@ -0,0 +1,2 @@
+*
+!.gitignore

data/.gitignore

@@ -0,0 +1,48 @@
+from datasets.utils import SeededDataLoader
+import torchvision.transforms as T
+from torchvision.datasets import CIFAR10
+from globals import CONFIG
+
+def get_transform(size, padding, mean, std, preprocess):
+    transform = []
+    transform.append(T.Resize((size, size)))
+    if preprocess:
+        transform.append(T.RandomCrop(size=size, padding=padding))
+        transform.append(T.RandomHorizontalFlip())
+    transform.append(T.ToTensor())
+    transform.append(T.Normalize(mean, std))
+    return T.Compose(transform)
+
+def load_data():
+    size = 32
+    if 'pretrain' in CONFIG.experiment_args:
+        size = 224
+
+    CONFIG.num_classes = 10
+    CONFIG.data_input_size = (3, size, size)
+
+    mean, std = (0.491, 0.482, 0.447), (0.247, 0.243, 0.262)
+    train_transform = get_transform(size=size, padding=4, mean=mean, std=std, preprocess=True)
+    test_transform = get_transform(size=size, padding=4, mean=mean, std=std, preprocess=False)
+    train_dataset = CIFAR10(CONFIG.dataset_args['root'], train=True, download=True, transform=train_transform)
+    test_dataset = CIFAR10(CONFIG.dataset_args['root'], train=False, download=True, transform=test_transform)
+
+    train_loader = SeededDataLoader(
+        train_dataset,
+        batch_size=CONFIG.batch_size,
+        shuffle=True,
+        num_workers=CONFIG.num_workers,
+        pin_memory=True,
+        persistent_workers=True
+    )
+
+    test_loader = SeededDataLoader(
+        test_dataset,
+        batch_size=CONFIG.batch_size,
+        shuffle=False,
+        num_workers=CONFIG.num_workers,
+        pin_memory=True,
+        persistent_workers=True
+    )
+
+    return {'train': train_loader, 'test': test_loader}

datasets/CIFAR10/dataset.py

@@ -0,0 +1,48 @@
+from datasets.utils import SeededDataLoader
+import torchvision.transforms as T
+from torchvision.datasets import CIFAR100
+from globals import CONFIG
+
+def get_transform(size, padding, mean, std, preprocess):
+    transform = []
+    transform.append(T.Resize((size, size)))
+    if preprocess:
+        transform.append(T.RandomCrop(size=size, padding=padding))
+        transform.append(T.RandomHorizontalFlip())
+    transform.append(T.ToTensor())
+    transform.append(T.Normalize(mean, std))
+    return T.Compose(transform)
+
+def load_data():
+    size = 32
+    if 'pretrain' in CONFIG.experiment_args:
+        size = 224
+
+    CONFIG.num_classes = 100
+    CONFIG.data_input_size = (3, size, size)
+
+    mean, std = (0.507, 0.487, 0.441), (0.267, 0.256, 0.276)
+    train_transform = get_transform(size=size, padding=4, mean=mean, std=std, preprocess=True)
+    test_transform = get_transform(size=size, padding=4, mean=mean, std=std, preprocess=False)
+    train_dataset = CIFAR100(CONFIG.dataset_args['root'], train=True, download=True, transform=train_transform)
+    test_dataset = CIFAR100(CONFIG.dataset_args['root'], train=False, download=True, transform=test_transform)
+
+    train_loader = SeededDataLoader(
+        train_dataset,
+        batch_size=CONFIG.batch_size,
+        shuffle=True,
+        num_workers=CONFIG.num_workers,
+        pin_memory=True,
+        persistent_workers=True
+    )
+
+    test_loader = SeededDataLoader(
+        test_dataset,
+        batch_size=CONFIG.batch_size,
+        shuffle=False,
+        num_workers=CONFIG.num_workers,
+        pin_memory=True,
+        persistent_workers=True
+    )
+
+    return {'train': train_loader, 'test': test_loader}

datasets/CIFAR100/dataset.py

@@ -0,0 +1,57 @@
+import os
+from datasets.utils import SeededDataLoader
+import torchvision.transforms as T
+
+from torchvision import datasets
+
+from globals import CONFIG
+
+def get_transform(train, mean, std):
+    if train:
+        transform = T.Compose([
+            T.RandomResizedCrop(224, scale=(0.2,1.)),
+            T.RandomGrayscale(p=0.2),
+            T.ColorJitter(0.4, 0.4, 0.4, 0.4),
+            T.RandomHorizontalFlip(),
+            T.ToTensor(),
+            T.Normalize(mean, std)])
+    else:
+        transform = T.Compose([
+            T.Resize(256),
+            T.CenterCrop(224),
+            T.ToTensor(),
+            T.Normalize(mean, std)])
+    return transform
+
+def load_data():
+    CONFIG.num_classes = 1000
+    CONFIG.data_input_size = (3, 224, 224)
+
+    mean, std = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
+    train_transform = get_transform(train=True, mean=mean, std=std)
+    test_transform = get_transform(train=False, mean=mean, std=std)
+
+    train_folder = os.path.join(CONFIG.dataset_args['root'], 'train')
+    train_dataset = datasets.ImageFolder(train_folder, transform=train_transform)
+    test_folder = os.path.join(CONFIG.dataset_args['root'], 'val')
+    test_dataset = datasets.ImageFolder(test_folder, transform=test_transform)
+
+    train_loader = SeededDataLoader(
+        train_dataset,
+        batch_size=CONFIG.batch_size,
+        shuffle=True,
+        num_workers=CONFIG.num_workers,
+        pin_memory=True,
+        persistent_workers=True
+    )
+
+    test_loader = SeededDataLoader(
+        test_dataset,
+        batch_size=CONFIG.batch_size,
+        shuffle=False,
+        num_workers=CONFIG.num_workers,
+        pin_memory=True,
+        persistent_workers=True
+    )
+
+    return {'train': train_loader, 'test': test_loader}

datasets/ImageNet/dataset.py

@@ -0,0 +1,105 @@
+import io
+import pandas as pd
+import glob
+import os
+from shutil import move
+from os.path import join
+from os import listdir, rmdir
+from datasets.utils import SeededDataLoader
+import torchvision.transforms as T
+
+from torchvision import datasets
+
+from globals import CONFIG
+
+# Based on https://github.com/tjmoon0104/pytorch-tiny-imagenet/blob/master/val_format.py
+def TINYIMAGENET(root, train=True, transform=None, target_transform=None, download=False):
+    
+    def _exists(root, filename):
+        return os.path.exists(os.path.join(root, filename))
+
+    def _download(url, root, filename):
+        datasets.utils.download_and_extract_archive(url=url, 
+                                                    download_root=root, 
+                                                    extract_root=root, 
+                                                    filename=filename)
+
+    def _setup(root, base_folder):
+        target_folder = os.path.join(root, base_folder, 'val/')
+
+        val_dict = {}
+        with open(target_folder + 'val_annotations.txt', 'r') as f:
+            for line in f.readlines():
+                split_line = line.split('\t')
+                val_dict[split_line[0]] = split_line[1]
+                
+        paths = glob.glob(target_folder + 'images/*')
+        paths[0].split('/')[-1]
+        for path in paths:
+            file = path.split('/')[-1]
+            folder = val_dict[file]
+            if not os.path.exists(target_folder + str(folder)):
+                os.mkdir(target_folder + str(folder))
+                
+        for path in paths:
+            file = path.split('/')[-1]
+            folder = val_dict[file]
+            dest = target_folder + str(folder) + '/' + str(file)
+            move(path, dest)
+            
+        os.remove(target_folder + 'val_annotations.txt')
+        rmdir(target_folder + 'images')
+
+    url = "http://cs231n.stanford.edu/tiny-imagenet-200.zip"
+    filename = "tiny-imagenet-200.zip"
+    base_folder = 'tiny-imagenet-200'
+
+    if download and not _exists(root, filename):
+        _download(url, root, filename)
+        _setup(root, base_folder)
+    folder = os.path.join(root, base_folder, 'train' if train else 'val')
+
+    return datasets.ImageFolder(folder, transform=transform, target_transform=target_transform)
+
+def get_transform(size, mean, std, preprocess):
+    transform = []
+    if preprocess:
+        transform.append(T.RandomResizedCrop(size=size, scale=(0.1, 1.0), ratio=(0.8, 1.25)))
+        transform.append(T.RandomHorizontalFlip())
+    transform.append(T.ToTensor())
+    transform.append(T.Normalize(mean, std))
+    return T.Compose(transform)
+
+def load_data():
+    size = 64
+    if 'pretrain' in CONFIG.experiment_args:
+        size = 224
+
+    CONFIG.num_classes = 200
+    CONFIG.data_input_size = (3, size, size)
+
+    mean, std = (0.480, 0.448, 0.397), (0.276, 0.269, 0.282)
+    train_transform = get_transform(size=size, mean=mean, std=std, preprocess=True)
+    test_transform = get_transform(size=size, mean=mean, std=std, preprocess=False)
+    train_dataset = TINYIMAGENET(CONFIG.dataset_args['root'], train=True, download=True, transform=train_transform)
+    test_dataset = TINYIMAGENET(CONFIG.dataset_args['root'], train=False, download=True, transform=test_transform)
+
+    train_loader = SeededDataLoader(
+        train_dataset,
+        batch_size=CONFIG.batch_size,
+        shuffle=True,
+        num_workers=CONFIG.num_workers,
+        pin_memory=True,
+        persistent_workers=True
+    )
+
+    test_loader = SeededDataLoader(
+        test_dataset,
+        batch_size=CONFIG.batch_size,
+        shuffle=False,
+        num_workers=CONFIG.num_workers,
+        pin_memory=True,
+        persistent_workers=True
+    )
+
+    return {'train': train_loader, 'test': test_loader}