Modelpublic

lalodeavilaarm/spiral-vit-medmnist-ablation

sourceHugging Facemitupdated 10d agoView on Hugging Face

0likes

Model Card

SPIRAL-ViT — MedMNIST Full-Resolution Ablation Checkpoints

Full checkpoint set (8 configurations × 12 MedMNIST v2 datasets, 96 files, native 224×224 resolution) for the foveated spectral-spatial Vision Transformer ablation study — spiral token ordering + learned-DCT frequency-domain patch embedding + Gaussian attention bias.

Training/evaluation code: github.com/Tamez-Research-Group/FoveatedSpectral-ViT

Ablation configuration map

Config	token_order	embed_type	gaussian_bias
C1 (baseline)	raster	conv	False
C2	spiral	conv	False
C3	raster	learned_dct	False
C4	raster	conv	True
C5	spiral	learned_dct	False
C6	spiral	conv	True
C7	raster	learned_dct	True
C8 (full model)	spiral	learned_dct	True

Repository layout

<dataset>/vit_<dataset>_best_<embed_type>_<token_order>_<gauss_tag>.pth

One folder per dataset (blood, breast, chest, derma, oct, organa, organc, organs, path, pneumonia, retina, tissue), 8 checkpoints each — one per config above. <gauss_tag> is gaussian or nogaussian.

Each checkpoint is a dict:

python

{"state_dict": ..., "cfg": {...}, "epoch": int, "best_f1": float}

cfg records the exact architecture (embed_type, token_order, gaussian_bias, image_size, patch_size, dim, depth, heads, mlp_dim, dropout, freq_keep) needed to reconstruct the model.

Loading a checkpoint

Requires the ViT class from the code repo above (src/models/vit.py):

python

import torch
from huggingface_hub import hf_hub_download
from models.vit import ViT  # from FoveatedSpectral-ViT/src

NUM_CLASSES = {  # MedMNIST v2 label counts
    "blood": 8, "breast": 2, "chest": 14, "derma": 7, "oct": 4,
    "organa": 11, "organc": 11, "organs": 11, "path": 9,
    "pneumonia": 2, "retina": 5, "tissue": 8,
}

dataset, config_file = "blood", "vit_blood_best_learned_dct_spiral_gaussian.pth"  # C8
ckpt_path = hf_hub_download(repo_id="<your-username>/spiral-vit-medmnist-ablation",
                             filename=f"{dataset}/{config_file}")

ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
cfg = ckpt["cfg"]
model = ViT(
    image_size=int(cfg.get("image_size", 224)), patch_size=int(cfg.get("patch_size", 14)),
    num_classes=NUM_CLASSES[dataset], dim=int(cfg.get("dim", 256)), depth=int(cfg.get("depth", 6)),
    heads=int(cfg.get("heads", 8)), mlp_dim=int(cfg.get("mlp_dim", 512)), dropout=float(cfg.get("dropout", 0.1)),
    embed_type=str(cfg.get("embed_type", "conv")), token_order=str(cfg.get("token_order", "raster")),
    gaussian_bias=bool(cfg.get("gaussian_bias", False)), freq_keep=int(cfg.get("freq_keep", 16)),
)
model.load_state_dict(ckpt["state_dict"])
model.eval()

Checkpoint selection criterion: validation macro-F1 (not accuracy) for every config and dataset, including Chest (its 14-label multi-hot macro-F1, computed the same way as the rest).

License

MIT.