LSSB: Low-Shot learning with Shape Bias - CVPR2021

This is a repository containing PyTorch code for Using Shape to Categorize: Low-Shot Learning with an Explicit Shape Bias. You can find instructions for for testing pre-trained models in the paper and training your own models in the instructions below.

If you use our work in your research, please consider citing

@article{stojanov21cvpr,
      title={Using Shape to Categorize: Low-Shot Learning with an Explicit Shape Bias},
      author={Stefan Stojanov and Anh Thai and James M. Rehg},
      booktitle = {CVPR},
      year      = {2021}
}

Project Members

• Stefan Stojanov - Georgia Institute of Technology (PhD Student)
• Anh Thai - Georgia Institute of Technology (PhD Student)
• James M. Rehg - Georgia Institute of Technology

Contact

If you have any questions regarding the paper, data or code, please email Stefan Stojanov at [email protected]

Toys4K 3D Object Dataset

Details for downloading and rendering the Toys4K 3D object dataset are available in the Toys4K directory of this repository.

Results - Shape Bias Improves Low-Shot Generalization

The following results are averaged over three runs for FEAT and five runs for SimpleShot

ModelNet

ShapeNet

Toys4K

Installation

For this project we use miniconda to manage dependencies. After setting it up, you can install the lssb environment

conda env create -f environment.yml
conda activate lssb

pip install -e .

The most important dependencies to get right are PyTorch 1.6 and PyTorch Lightning 0.8.4. Since this project uses PyTorch Lightning, if you want to make significant modifications to the code it will be helpful to look into the overall PL workflow.

Make sure you install PyTorch with a compatible CUDA version. This repository has been tested with CUDA>10.1.

Datasets

You can download the data we used for the experiments for ModelNet, ShapeNet and our Toys4K dataset using download_data.sh.

Attribution information about our new dataset is available here. Please do not distribute the data made available here in any way, and please make sure you follow the respective ModelNet and ShapeNet licenses.

The approximate sizes for the datasets are 7.4GB for ModelNet, 11GB for ShapeNet55 and 4.1GB for Toys4K

For training and testing with shape bias use download_features.sh to download extracted DGCNN point cloud features.

Code Organization

Following is an outline of the LSSB codebase

lssb - main code directory
    - data - PyTorch dataset classes
        - modelnet.py - dataset class for ModelNet
        - sampler.py - low-shot batch sampler code
        - shapenet.py - datset class for ShapeNet
        - toys.py - dataset clas for Toys4K
        - transform.py - code for image and point cloud transformations
    - feat_extract - code for feature extraction from point clouds
        - extract_simpleshot_pc_feature_modelnet.py - ModelNet feature extraction
        - extract_simpleshot_pc_feature_shapenet.py - ShapeNet feature extraction
        - extract_simpleshot_pc_feature_toys4.5k.py - Toys feature extraction
        - utils.py - util functions for feature extraction
    - lowshot
        - configs
            - feat - config files for training FEAT models on all 3 datasets
            - simpleshot - config files for training SimpleShot models or all 3 datasets
        - models
            - base_feat_classifier.py - base class for FEAT classifier
            - base_simpleshot_classifier.py - base class for SimpleShot classifier
            - image_feat_classifier.py - subclass for image-based FEAT classifier
            - image_simpleshot_classifier.py - sublcass for image-based SimpleShot classifier 
            - joint_feat_classifier.py - subclass for shape-biased FEAT classifier
            - joint_simpleshot_classifier.py - subclass for shape-biased SimpleShot classifier
            - ptcld_simpleshot_classifier.py - subclass for point-cloud based SimpleShot classifier
        - utils
            - feat_utils.py - util functions specific to FEAT
            - simpleshot_utils.py - util functions specifict to SimpleShot
            - train_utils.py - just contains a function to return dataset objects; used across all methods
        - test.py - main testing script; used for all methods
        - train.py - main training script; used for all methods
    - nets
        - ResNet.py - borrowed from torchvision, definitons of ResNet models
        - dgcnn.py - definition of DGCNN point cloud classifier
        - feat.py - definition of FEAT transformer module

Evaluate Pre-Trained Models

First download the pretrained models (approx. 8.3GB)

bash download_models.sh

Once the models are downloaded, please use the scripts in testing_scripts to test the pretrained models e.g.

bash testing_scripts/test_simpleshot_modelnet.sh

Training From Scratch

To train models from scratch, please use the scripts in training_scripts e.g.

bash training_scripts/train_resnet18_simpleshot.sh 0 

Notes for training and testing with shape bias:

To train the shape biased image encoder, feature extraction needs to be done to save on compute. An example command is

python lssb/feat_extract/extract_simpleshot_pc_feature_modelnet.py --ckpt_path=pretrained_models/simpleshot/modelnet/ptcld-only/dgcnn-modelnet-simpleshot/version_0/checkpoints/epoch\=176_val_acc\=0.765.ckpt

Once this is completed, the .npz file should be moved to the dataset directory, under a new directory called features. You can download extracted features using download_features.sh as in the instructions above.

To train the shape-biased FEAT models, a shape-biased simpleshot checkpoint is needed. In the appropriate config file for the encoder add for example

encoder_path_resnet: pretrained_models/simpleshot/modelnet/shape-biased/joint-modelnet-pairwise-simpleshot/version_0/checkpoints/epoch=360_val_acc=0.718.ckpt

if you have downloaded the pretrained models.

Code Credits

We are grateful for the authors of the following code repositories, from where we borrow components used in this repository

• DGCNN
• FEAT
• PointNet/PointNet2-PyTorch
• PointNet2-PyTorch
• RFS
• SimpleShot