BatchFormer: Learning to Explore Sample Relationships for Robust Representation Learning

Introduction

This is the official PyTorch implementation of BatchFormer for Long-Tailed Recognition, Domain Generalization, Compositional Zero-Shot Learning, Contrastive Learning.

Sample Relationship Exploration for Robust Representation Learning

Please also refer to BatchFormerV2, in which we introduce a BatchFormerV2 module for vision Transformers.

Main Results

Long-Tailed Recognition

ImageNet-LT

Here, we demonstrate the result on one-stage RIDE (ResNext-50)

iNaturalist 2018

Object Detection (V2)

The backbone is ResNet-50. The training epoch is 50.

Panoptic segmentation (V2)

Contrastive Learning

Here, we provide the pretrained MoCo-V3 model corresponding to this strategy.

Domain Generalization

ResNet-18

Compositional Zero-Shot Learning

Few-Shot Learning

Experiments on CUB.

Image Classification (V2)

Reference

1. Long-tailed recognition by routing diverse distribution-aware experts. In ICLR, 2021
2. Parametric contrastive learning. In ICCV, 2021
3. Improved baselines with momentum contrastive learning.
4. An empirical study of training self-supervised vision transformers. In CVPR, 2021
5. Domain generalization by seeking flat minima. In NeurIPS, 2021.
6. Learning graph embeddings for compositional zero-shot learning. In CVPR, 2021
7. Contrastive learning based hybrid networks for long- tailed image classification. In CVPR, 2021

PyTorch Code

The proposed BatchFormer can be implemented with a few lines as follows,

def BatchFormer(x, y, encoder, is_training):
    # x: input features with the shape [N, C]
    # encoder: TransformerEncoderLayer(C,4,C,0.5)
    if not is_training:
        return x, y
    pre_x = x
    x = encoder(x.unsqueeze(1)).squeeze(1)
    x = torch.cat([pre_x, x], dim=0)
    y = torch.cat([y, y], dim=0)
    return x, y

Citation

If you find this repository helpful, please consider cite:

@inproceedings{hou2022batch,
    title={BatchFormer: Learning to Explore Sample Relationships for Robust Representation Learning},
    author={Hou, Zhi and Yu, Baosheng and Tao, Dacheng},
    booktitle={CVPR},
    year={2022}
}
@article{hou2022batchformerv2,
   title={BatchFormerV2: Exploring Sample Relationships for Dense Representation Learning},
   author={Hou, Zhi and Yu, Baosheng and Wang, Chaoyue and Zhan, Yibing and Tao, Dacheng},
   journal={arXiv preprint arXiv:2204.01254},
   year={2022}
}

Feel free to contact "zhou9878 at uni dot sydney dot edu dot au" if you have any questions.