wutong16 / DistributionBalancedLoss

[ ECCV 2020 Spotlight ] Pytorch implementation for "Distribution-Balanced Loss for Multi-Label Classification in Long-Tailed Datasets"
362 stars 46 forks source link

Distribution-Balanced Loss

[Paper]

The implementation of our paper Distribution-Balanced Loss for Multi-Label Classification in Long-Tailed Datasets (ECCV2020 Spotlight).

Tong WuQingqiu HuangZiwei LiuYu WangDahua Lin

Requirements

Installation

git clone git@github.com:wutong16/DistributionBalancedLoss.git
cd DistributionBalancedLoss

Quick start

Training

COCO-MLT

python tools/train.py configs/coco/LT_resnet50_pfc_DB.py 

VOC-MLT

python tools/train.py configs/voc/LT_resnet50_pfc_DB.py 

Testing

COCO-MLT

bash tools/dist_test.sh configs/coco/LT_resnet50_pfc_DB.py work_dirs/LT_coco_resnet50_pfc_DB/epoch_8.pth 1

VOC-MLT

bash tools/dist_test.sh configs/voc/LT_resnet50_pfc_DB.py work_dirs/LT_voc_resnet50_pfc_DB/epoch_8.pth 1

Pre-trained models

COCO-MLT

Backbone Total Head Medium Tail Download
ResNet-50 53.55 51.13 57.05 51.06 model

VOC-MLT

Backbone Total Head Medium Tail Download
ResNet-50 78.94 73.22 84.18 79.30 model

Datasets

Use our dataset

The long-tail multi-label datasets we use in the paper are created from MS COCO 2017 and Pascal VOC 2012. Annotations and statistics data resuired when training are saved under ./appendix in this repo.

appendix
  |--coco
    |--longtail2017
      |--class_freq.pkl
      |--class_split.pkl
      |--img_id.pkl
  |--VOCdevkit
    |--longtail2012
      |--class_freq.pkl
      |--class_split.pkl
      |--img_id.pkl

Try your own

You can also create a new long-tailed dataset by downloading the annotations, terse_gt_2017.pkl for COCO and terse_gt_2012.pkl for VOC, from here and move them into the right folders as below.

appendix
  |--coco
    |--longtail2017
      |--terse_gt_2017.pkl
  |--VOCdevkit
    |--longtail2012
      |--terse_gt_2012.pkl

Then run the following command, adjust the parameters as you like to control the distribution.

python tools/create_longtail_dataset.py

To update the corresponding class_freq.pkl files, please refer to def _save_info in .\mllt\datasets\custom.py.

License and Citation

The use of this software is RESTRICTED to non-commercial research and educational purposes.

@inproceedings{DistributionBalancedLoss,
  title={Distribution-Balanced Loss for Multi-Label Classification in Long-Tailed Datasets},
  author={Wu, Tong and Huang, Qingqiu and Liu, Ziwei and Wang, Yu and Lin, Dahua},
  booktitle={European Conference on Computer Vision (ECCV)},
  year={2020}
}

TODO

Contact

This repo is currently maintained by @wutong16 and @hqqasw