/ConstellationNet

(ICLR 2021) ConstellationNet: Attentional Constellation Nets for Few-Shot Learning

Primary LanguagePythonApache License 2.0Apache-2.0

Attentional Constellation Nets For Few-shot Learning

Introduction

This repository contains the official code and pretrained models for Attentional Constellation Nets For Few-shot Learning (ICLR 2021). In this paper, we tackle the few-shot learning problem and make an effort to enhance structured features by expanding CNNs with a constellation model, which performs cell feature clustering and encoding with a dense part representation; the relationships among the cell features are further modeled by an attention mechanism. With the additional constellation branch to increase the awareness of object parts, our method is able to attain the advantages of the CNNs while making the overall internal representations more robust in the few-shot learning setting. Our approach attains a significant improvement over the existing methods in few-shot learning on the CIFAR-FS, FC100, and mini-ImageNet benchmarks.

For more details, please refer to Attentional Constellation Nets For Few-shot Learning by Weijian Xu*, Yifan Xu*, Huaijin Wang*, and Zhuowen Tu.

Changelog

09/17/2021: Code for ConstellationNet are released.

Usage

Environment Preparation

  1. Set up a new conda environment and activate it.

    # Create an environment with Python 3.8.
conda create -n constells python==3.8
conda activate constells

  2. Install required packages.

    # Install PyTorch 1.8.0 w/ CUDA 11.1.
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge

# Install yaml
conda install -c anaconda pyyaml

# Install tensorboardx.
conda install -c conda-forge tensorboardx tqdm

Code and Datasets Preparation

  1. Clone the repo.

    git clone https://github.com/mlpc-ucsd/ConstellationNet.git
cd ConstellationNet

  2. Download datasets

 materials
├── mini-imagenet
│   ├── miniImageNet_category_split_test.pickle
│   ├── miniImageNet_category_split_train_phase_test.pickle
│   ├── miniImageNet_category_split_train_phase_train.pickle
│   ├── miniImageNet_category_split_train_phase_val.pickle
│   ├── miniImageNet_category_split_val.pickle
├── cifar-fs
│   ├── CIFAR_FS_test.pickle
│   ├── CIFAR_FS_train.pickle
│   ├── CIFAR_FS_val.pickle
├── fc100
│   ├── FC100_test.pickle
│   ├── FC100_train.pickle
│   ├── FC100_val.pickle

Pre-trained Checkpoints

We provide the Constellation Nets checkpoints pre-trained on the Mini-Imagenet, CIFAR-FS and FC100.

Dataset Model Backbone Acc@ 5-way 1-shot Acc@ 5-way 5-shot #Params SHA-256 (first 8 chars) URL
Mini-ImageNet ConstellationNets Conv-4 59.67 ± 0.23 75.98 ± 0.18 200K d76075a5 model
Mini-ImageNet ConstellationNets ResNet-12 65.53 ± 0.23 80.55 ± 0.16 8.4M cf716d90 model
CIFAR-FS ConstellationNets Conv-4 69.1 ± 0.3 83.0 ± 0.2 200K 4ea590f9 model
CIFAR-FS ConstellationNets ResNet-12 76.1 ± 0.2 87.4 ± 0.2 8.4M dc5d56fa model
FC100 ConstellationNets ResNet-12 43.9 ± 0.2 59.7 ± 0.2 8.4M d9a829f7 model

Train

The following commands provide an example to train the Constellation Net .

# Usage: bash ./scripts/train.sh [Dataset (mini, cifar-fs, fc100)] [Backbone (conv4, res12)] [GPU index] [Tag]
bash ./scripts/train.sh mini conv4 0 trial1

Evaluate

The following commands provide an example to evaluate the checkpoint after training.

# Usage: bash ./scripts/test.sh [Dataset (mini, cifar-fs, fc100)] [Backbone (conv4, res12)] [GPU index] [Tag]
bash ./scripts/eval.sh mini conv4 0 trial1

Citation

@inproceedings{
xu2021attentional,
title={Attentional Constellation Nets for Few-Shot Learning},
author={Weijian Xu and yifan xu and Huaijin Wang and Zhuowen Tu},
booktitle={International Conference on Learning Representations},
year={2021},
url={https://openreview.net/forum?id=vujTf_I8Kmc}
}

License

This repository is released under the Apache License 2.0. License can be found in LICENSE file.

Acknowledgment