Superclass-Conditional Gaussian Mixture Model For Learning Fine-Grained Embeddings (ICLR 2022 Spotlight)
This is the code for the paper "Superclass-Conditional Gaussian Mixture Model for Learning Fine-Grained Embeddings" in ICLR 2022 (pdf). This code provides a demo on BREEDS dataset, and it can be adapted to other datasets including CIFAR-100 and tieredImageNet.
The experiments were done using python3.7, with the following packages:
- learn2learn==0.1.5
- matplotlib==3.4.2
- networkx==2.5.1
- numpy==1.20.3
- pandas==1.3.0
- robustness==1.2.1.post2
- scikit-learn==0.24.2
- scipy==1.7.0
- seaborn==0.11.1
- torch==1.4.0+cu92
- torchvision==0.5.0+cu92
- Download the ImageNet dataset.
- Following the official BREEDS repo, run
import os from robustness.tools.breeds_helpers import setup_breeds info_dir= "[your_imagenet_path]/ILSVRC/BREEDS" if not (os.path.exists(info_dir) and len(os.listdir(info_dir))): print("Downloading class hierarchy information into `info_dir`") setup_breeds(info_dir)- The directory structure is
└── ILSVRC ├── Annotations │ └── CLS-LOC ├── BREEDS │ ├── class_hierarchy.txt │ ├── dataset_class_info.json │ └── node_names.txt ├── Data │ └── CLS-LOC ├── ImageSets │ └── CLS-LOC └── Meta ├── imagenet_class_index.json ├── test.json ├── wordnet.is_a.txt └── words.txt CIFAR-100 can be downloaded from [link].
- Once downloaded, use
dataset_cifar.pyindataset/folder to generate minibatches for model training.
TieredImageNet can be downloaded from [link].
- Once downloaded, use
dataset_tiered_imagenet.pyindataset/folder to generate minibatches for model training.
First, create a directory to save the pre-trained models.
mkdir pretrain_model To train SCGM with a generic encoder (i.e., SCGM-G) on Living17 dataset, run
python train_scgm_g.py \ --data [path to data directory] \ --workers 32 \ --epochs 200 \ --batch_size 256 \ --hiddim 128 \ --tau 0.1 \ --alpha 0.5 \ --lmd 25 \ --n-subclass 100 \ --n-class 17 \ --dataset living17 To train SCGM with a momentum-based encoder (i.e., SCGM-A) on Living17 dataset, run
python train_scgm_g.py \ --data [path to data directory] \ --arch resnet50 \ --workers 32 \ --epochs 200 \ --batch_size 256 \ --hiddim 128 \ --queue multi \ --metric angular \ --head-type seq_em \ --cst-t 0.2 \ --tau1 0.1 \ --alpha 0.5 \ --lmd 25 \ --n-subclass 100 \ --n-class 17 \ --dataset living17 The default parameters were set for training on BREEDS dataset. To check the model parameters, run
python train_scgm_g.py -h python train_scgm_a.py -h To test the performance of the pre-trained SCGM-G on the cross-granularity few-shot (CGFS) learning setting, run
python test_scgm_g.py --data [path to data directory] \ --batch_size 256 \ --n-test-runs 1000 \ --n-ways 5 \ --n-shots 1 \ --n-queries 15 \ --feat-norm \ --classifier LR \ --hiddim 128 \ --n-subclass 100 \ --n-class 17 \ --dataset living17 To test the performance of the pre-trained SCGM-A, run
python test_scgm_a.py --data [path to data directory] \ --arch resnet50 \ --batch_size 256 \ --n-test-runs 1000 \ --n-ways 5 \ --n-shots 1 \ --n-queries 15 \ --feat-norm \ --classifier LR \ --hiddim 128 \ --n-subclass 100 \ --n-class 17 \ --dataset living17 Similarly, to test the performance on the fine-grained intra-class setting, run
python test_fg_scgm_g.py python test_fg_scgm_a.py To visualize the embeddings, include the lines 318 to 340 in train_scgm_g.py and the lines 341 to 363 in train_scgm_a.py.
@inproceedings{ni2022superclass, title={Superclass-Conditional Gaussian Mixture Model For Learning Fine-Grained Embeddings}, author={Ni, Jingchao and Cheng, Wei and Chen, Zhengzhang and Asakura, Takayoshi and Soma, Tomoya and Kato, Sho and Chen, Haifeng}, booktitle={International Conference on Learning Representations}, year={2022} }