Metadata-Version: 2.4
Name: wsl-modelzoo
Version: 0.1.0
Summary: CLI tool for downloading checkpoints and metadata developed by the AIML Lab.
Author-email: Salvatore Pascarella <salvatore.pascarella01@gmail.com>, Leo Meynent <leo.meynent@unisg.ch>
License: MIT
Keywords: model-zoo,weight-space-learning,deep-learning,machine-learning,artificial-intelligence
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: requests>=2.31
Requires-Dist: tqdm>=4.66

# Model Zoo Downloader

A command-line tool for downloading checkpoints and metadata developed by the AIML Lab.

This package provides a simple CLI to fetch artifacts from the **Core**, **EuroSAT**, and **ViT** model zoos.

## Installation

Install the package:

```bash
pip install wsl-modelzoo
```
This tool currently also requires the `wget` executable to be available on your system.
- macOS: `brew install wget`
- Ubuntu/Debian: `sudo apt install wget`

## Model Zoos
| Model Zoo | Available | Paper
| :--------: | :-------: | :-------:
| Core |✅| [Link](https://arxiv.org/abs/2209.14764)
| EuroSAT |✅| [Link](https://ieeexplore.ieee.org/iel7/10281394/10281399/10283060.pdf)
| ViT |✅| [Link](https://arxiv.org/abs/2504.10231)
| Phase Transitions|❌| [Link](https://arxiv.org/abs/2504.18072)

## Download directory

The CLI supports a `--dir` argument that lets you specify the root directory where checkpoints and metadata will be downloaded.

If `--dir` is not provided, the tool creates a `downloads` directory in the package root and stores all downloaded files there.

## Core Model Zoo 

### Arguments

| Argument | Required | Description                |
| :--------: | :-------: | :------------------------- |
| --zoo |✅|Select the model zoo (`core`)
| --arch | ✅| Model architecture
| --dataset | ✅| Dataset 
| --config| ✅|Configuration 
| --seed | ❌| Model seed(s)
| --ckpts | ❌| Checkpoints indices to fetch|
| --dropout |  ❌|Dropout
| --init| ❌|Initialization Type
| --activation| ❌|Activation Function
| --lr| ❌|Learning Rate
| --optimizer| ❌|Optimizer
| --wd| ❌|Weight Decay

### Core Model Zoo: valid combinations of architectures and datasets
Datasets marked with `*` also have a sparsified variant, for example `mnist_sparsified`.

| Architecture | Dataset | 
| :-------- | :------- | 
|`cnn-small`|`mnist*`, `fmnist*`, `svhn*`, <br>`usps*`, `cifar10*`, `stl10*`|
|`cnn-large`|`cifar10*`, `stl10*`| 
|`resnet18`|`cifar10`, `cifar100`, <br>`tinyimagenet`| 

### Valid config combinations

| Architecture | Dataset | Supported configs |
|---|---|---|
| `cnn-small` | `mnist*`, `fmnist*`, `svhn*`, `usps`, `cifar10`, `stl10` | `uniform`, `fixed`, `random` |
| `cnn-small` | `usps_sparsified`, `cifar10_sparsified`, `stl10_sparsified` | `uniform` |
| `cnn-large` | `cifar10*`, `stl10*`  | `uniform`, `fixed`, `random` |
| `resnet18` | `cifar10`, `cifar100`, `tinyimagenet`  | `uniform` |

### Extra hyperparameter filters

These filters only apply to selected `cnn-small` and `cnn-large` combinations with `fixed` or `random` configs.

Currently, extra filters accept only one value per argument unlike `--seed` and `--ckpts` (see below)

| Architecture | Dataset | Config | Supported extra filters |
|---|---|---|---|
| `cnn-small` | `mnist`, `fmnist`, `svhn` <br> `usps`, `cifar10`, `stl10` | `fixed`, `random` | `--dropout`, `--init`, `--activation`, <br> `--lr`, `--optimizer`, `--wd` |
| `cnn-large` | `cifar10`, `stl10`  | `fixed`, `random` | `--dropout`, `--init`, `--activation`, <br> `--lr`, `--optimizer`, `--wd` |

### Allowed values for extra hyperparameter filters

| Filter | Applies to | Allowed values |
|---|---|---|
| `--init` | all supported extra-filter combinations | `uniform`, `normal`, `kaiming_uniform`, `kaiming_normal` |
| `--activation` | all supported extra-filter combinations | `tanh`, `relu`, `gelu`, `sigmoid` |
| `--optimizer` | most supported extra-filter combinations | `adam`, `sgd` |
| `--optimizer` | `cnn-small` + `svhn` + (`fixed` or `random`) | `sgd` |
| `--dropout` | most supported extra-filter combinations | `0`, `0.5` |
| `--dropout` | `cnn-small` + `svhn` + (`fixed` or `random`) | `0`, `0.3`, `0.5` |
| `--lr` | most supported extra-filter combinations | `0.001`, `0.0001` |
| `--lr` | `cnn-small` + `svhn` + (`fixed` or `random`) | `0.1`, `0.001`, `0.0001` |
| `--lr` | `cnn-small` + `cifar10` + `random` | `0.001` |
| `--lr` | `cnn-large` + `cifar10` + (`fixed` or `random`) | `0.001` |
| `--wd` | most supported extra-filter combinations | `0.01`, `0.001` |
| `--wd` | `cnn-small` + `svhn` + (`fixed` or `random`) | `0`, `0.001`, `0.1` |

`--seed` cannot be used with `fixed` or `random` configs.

### Seed and checkpoint ranges

`--seed` and `--ckpts` accept one or more values within their allowed ranges.

You can pass:

- individual values, for example `1 10 25`
- inclusive ranges, for example `25-30`
- the special value `all` to download every available seed or checkpoint

By default, all the available seeds and checkpoints are retrieved if the user skips one of the two arguments above.

| Parameter | Range (inclusive) |
|---|---|
| `--seed` | `1..1000` |

| Architecture / dataset | Checkpoint range (inclusive) |
|---|---|
| `cnn-small` + `mnist` | `0..50` |
| `cnn-small` + `mnist_sparsified` | `0..25` |
| `cnn-small` + `fmnist` | `0..50` |
| `cnn-small` + `fmnist_sparsified` | `0..25`|
| `cnn-small` + `svhn` | `0..50` |
| `cnn-small` + `svhn_sparsified` | `0..25` |
| `cnn-small` + `usps` | `0..50` |
| `cnn-small` + `usps_sparsified` | `0..25` |
| `cnn-small` + `cifar10` | `0..50` |
| `cnn-small` + `cifar10_sparsified` | `0..25` |
| `cnn-small` + `stl10` | `0..50` |
| `cnn-small` + `stl10_sparsified` | `0..25` |
| `cnn-large` + `cifar10` | `0..50` |
| `cnn-large` + `cifar10_sparsified` | `0..25` |
| `cnn-large` + `stl10` | `0..50` |
| `cnn-large` + `stl10_sparsified` | `0..25` |
| `resnet18` + `cifar10` | `0..50` |
| `resnet18` + `cifar100` | `0..60` |
| `resnet18` + `tinyimagenet` | `0..60` |

### Examples
```bash
modelzoo fetch --zoo core --arch cnn-small --dataset mnist --config uniform --seed all --ckpts 0 10 20 30-45 --dir /your/destination/path

modelzoo fetch --zoo core --arch cnn-small --dataset mnist --config fixed --init kaiming_uniform --dropout 0.5 --activation relu --lr 0.001 --optim adam --wd 0.001 --ckpts 0 10 20

modelzoo fetch --zoo core --arch cnn-small --dataset mnist_sparsified --config uniform --ckpts 0 5 10

modelzoo fetch --zoo core --arch cnn-large --dataset stl10 --config fixed --init kaiming_normal --dropout 0.5 --activation tanh --lr 0.0001 --optim adam --wd 0.001 --ckpts 0 10 20

modelzoo fetch --zoo core --arch resnet18 --dataset cifar100 --config uniform --seed 42 --ckpts all
```

## EuroSAT Model Zoo

### Arguments

| Argument | Required | Description |
|---|:---:|---|
| `--zoo` | ✅ | Select the model zoo (`eurosat`) |
| `--arch` | ✅ | Model architecture |
| `--dataset` | ✅ | Dataset |
| `--config` | ❌ | Sparsified configuration |
| `--seed` | ❌ | Model seed(s) |
| `--ckpts` | ❌ | Checkpoint indices to fetch |

`--config` is required when `--dataset eurosat_sparsified` is used. <br>
`--seed` is not available when `--dataset eurosat_sparsified` is used. All the available seeds will be automatically downloaded

### Valid architecture / dataset combinations

| Architecture | Supported datasets |
|---|---|
| `cnn-small` | `eurosat`, `eurosat_sparsified` |

### Valid config combinations

Currently, `--config` accepts only a single value

| Dataset | Supported configs |
|---|---|
| `eurosat` | no extra config required |
| `eurosat_sparsified` | `ard_vd`, `magn_10`, `magn_20`, `magn_30`, `magn_40`, `magn_50`, `magn_60`, `magn_70`, `magn_80`, `magn_90` |

### Seed and checkpoint ranges

`--seed` and `--ckpts` accept one or more values within their allowed ranges.

You can pass:

- individual values, for example `1 10 25`
- ranges, for example `25-30`
- the special value `all` to download every available seed or checkpoint

By default, all the available seeds and checkpoints are retrieved if the user skips one of the two arguments above.

| Parameter | Range (inclusive) |
|---|---|
| `--seed` | `1..1000` |

| Dataset / config | Checkpoint range (inclusive) |
|---|---|
| `eurosat` | `0..50` |
| `eurosat_sparsified` + `ard_vd` | `0..25` |
| `eurosat_sparsified` + `magn_10` | `0..15` |
| `eurosat_sparsified` + `magn_20` | `0..15` |
| `eurosat_sparsified` + `magn_30` | `0..15` |
| `eurosat_sparsified` + `magn_40` | `0..15` |
| `eurosat_sparsified` + `magn_50` | `0..15` |
| `eurosat_sparsified` + `magn_60` | `0..15` |
| `eurosat_sparsified` + `magn_70` | `0..15` |
| `eurosat_sparsified` + `magn_80` | `0..15` |
| `eurosat_sparsified` + `magn_90` | `0..15` |

### Examples

```bash
modelzoo fetch --zoo eurosat --arch cnn-small --dataset eurosat --seed 1 --ckpts 0 5 10 --dir /your/destination/path

modelzoo fetch --zoo eurosat --arch cnn-small --dataset eurosat_sparsified --config ard_vd --ckpts 0-10

modelzoo fetch --zoo eurosat --arch cnn-small --dataset eurosat_sparsified --config magn_50 --ckpts all
```

## ViT Model Zoo

The ViT model zoo supports:

- **pretrained models**
- **pretrained + corresponding finetuned models**

To download **only pretrained models**, provide only the pretraining arguments and **omit all finetuning arguments**.

It is **not possible to download only finetuned models**. Finetuned models are always selected starting from a matching pretrained model.

### Arguments

#### Pretraining arguments

| Argument | Required | Description |
|---|:---:|---|
| `--zoo` | ✅ | Select the model zoo (`vit`) |
| `--dataset_pt` | ✅ | Pretraining dataset |
| `--config` | ✅ | Pretraining regime |
| `--seed_pt` | ❌ | Pretraining seed(s) |
| `--ckpts_pt` | ❌ | Pretraining checkpoint indices |
| `--ddp` | ❌ | Use the DDP supervised-pretraining variant (`true` or `false`) |

#### Finetuning arguments

| Argument | Required | Description |
|---|:---:|---|
| `--dataset_ft` | ❌ | Finetuning dataset |
| `--seed_ft` | ❌ | Finetuning seed(s) |
| `--lr_ft` | ❌ | Finetuning learning rate |
| `--optim_ft` | ❌ | Finetuning optimizer |
| `--hd_ft` | ❌ | Finetuning head hidden dimension |
| `--fcmlp_ft` | ❌ | Whether the finetuning head uses an MLP |
| `--ckpts_ft` | ❌ | Finetuning checkpoint indices |

### Valid pretraining combinations

| Argument | Allowed values |
|---|---|
| `--dataset_pt` | `imagenet` |
| `--config` | `sl`, `cl` |
| `--ddp` | `true`, `false` (default) |

### Valid pretraining seeds

| Pretraining mode | Allowed seeds |
|---|---|
| `sl` + `--ddp false` | `42`, `6312` |
| `sl` + `--ddp true` | `1957`, `2517`, `3753`, `6312`, `7853` |
| `cl` | `42`, `2517`, `6312` |

### Pretraining checkpoint range

| Parameter | Range (inclusive) |
|---|---|
| `--ckpts_pt` | `0..89` |

- If `--seed_pt` is omitted, all available pretraining seeds are used.
- If `--ckpts_pt` is omitted, all available pretraining checkpoints are used.

You can pass:

- individual values, for example `1 10 25`
- ranges, for example `25-30`
- the special value `all` to download every available seed or checkpoint

By default, all the available seeds and checkpoints are retrieved if the user skips one of the two arguments above.

### Valid finetuning combinations

Finetuning is only available **on top of a selected pretrained model**.

| Argument | Allowed values |
|---|---|
| `--dataset_ft` | `cifar100` |
| `--seed_ft` | `42`, `6312` |
| `--lr_ft` | `0.003`, `0.0001`, `0.001` |
| `--optim_ft` | `sgd`, `adamw` |
| `--hd_ft` | `256`, `none` |
| `--fcmlp_ft` | `true`, `false` |

### Finetuning checkpoint range

| Parameter | Range (inclusive) |
|---|---|
| `--ckpts_ft` | `0..49` |

### Important usage notes

| Use case | What to do |
|---|---|
| Download only pretrained models | Provide only `--dataset_pt`, `--config`, and optional pretraining arguments such as `--seed_pt`, `--ckpts_pt`, `--ddp` |
| Download pretrained and finetuned models | Provide the pretraining arguments **and** the finetuning arguments |
| Download only finetuned models | Not supported |


## Examples

```bash
modelzoo fetch --zoo vit --dataset_pt imagenet --config sl --ddp false --seed_pt 42 --ckpts_pt 0 10 20 --dir /your/destination/path

modelzoo fetch --zoo vit --dataset_pt imagenet --config sl --ddp true --seed_pt 1957 --ckpts_pt 0 10

modelzoo fetch --zoo vit --dataset_pt imagenet --config cl --ddp false --seed_pt 2517 --ckpts_pt 0 20

modelzoo fetch --zoo vit --dataset_pt imagenet --config sl --ddp false --seed_pt 42 --ckpts_pt 0 10 --dataset_ft cifar100 --seed_ft 42 --lr_ft 0.001 --optim_ft adamw --hd_ft 256 --fcmlp_ft true --ckpts_ft 0 10

modelzoo fetch --zoo vit --dataset_pt imagenet --config cl --ddp false --seed_pt 6312 --ckpts_pt 0 20 --dataset_ft cifar100 --seed_ft 6312 --lr_ft 0.003 --optim_ft sgd --hd_ft none --fcmlp_ft false --ckpts_ft 0 5 10
```

# Get in touch!
If there are any questions, reach out to us! We'll respond to issues here, or via email `salvatore.pascarella@unisg.ch`, `leo.meynent@unisg.ch`. 
