基于标签分布感知间隔损失的学习不平衡数据集

Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, Tengyu Ma

---

这是论文《基于标签分布感知间隔损失的学习不平衡数据集》（arXiv:1906.07413）中LDAM-DRW的官方PyTorch实现。

依赖项

代码基于以下库构建：

• PyTorch 1.2
• TensorboardX
• scikit-learn

数据集

• 不平衡的CIFAR。原始数据将通过imbalancec_cifar.py下载并转换。
• 论文还报告了在Tiny ImageNet和iNaturalist 2018上的实验结果。我们将在后续更新这些数据集的相关代码。

训练

本仓库提供了若干训练示例：

• 在长尾分布且比例为100的不平衡数据上训练ERM基线：

python cifar_train.py --gpu 0 --imb_type exp --imb_factor 0.01 --loss_type CE --train_rule None

• 在长尾分布且比例为100的不平衡数据上同时使用LDAM损失和DRW训练策略进行训练：

python cifar_train.py --gpu 0 --imb_type exp --imb_factor 0.01 --loss_type LDAM --train_rule DRW

参考文献

如果您觉得我们的论文和代码库有用，请引用如下：

@inproceedings{cao2019learning,
  title={Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss},
  author={Cao, Kaidi and Wei, Colin and Gaidon, Adrien and Arechiga, Nikos and Ma, Tengyu},
  booktitle={Advances in Neural Information Processing Systems},
  year={2019}
}

LDAM-DRW 快速上手指南

本指南帮助中国开发者快速部署并使用 LDAM-DRW（Label-Distribution-Aware Margin Loss），用于解决长尾分布下的不平衡数据集学习问题。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 建议 Python 3.6+
• 核心依赖:
  • PyTorch >= 1.2
  • TensorboardX
  • scikit-learn

国内加速建议：安装 PyTorch 时，推荐使用清华或中科大镜像源以提升下载速度。 例如使用 pip 安装： pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆项目代码

   git clone https://github.com/kaidic/LDAM-DRW.git
   cd LDAM-DRW
   

2. 安装依赖库 项目未提供 requirements.txt，请手动安装所需库（建议使用国内镜像）：

   pip install tensorboardX scikit-learn -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   注：PyTorch 请根据您的 CUDA 版本前往官网或上述镜像源单独安装。

3. 准备数据 本项目支持不平衡 CIFAR 数据集。首次运行时，脚本会自动下载原始数据并通过 imbalancec_cifar.py 进行转换处理，无需手动干预。

基本使用

以下命令展示了如何在长尾不平衡比例（Imbalance Ratio）为 100 的 CIFAR 数据集上进行训练。

1. 训练基准模型 (ERM Baseline)

使用标准的交叉熵损失函数（CE）作为对照：

python cifar_train.py --gpu 0 --imb_type exp --imb_factor 0.01 --loss_type CE --train_rule None

2. 训练 LDAM-DRW 模型 (推荐)

使用本文提出的 LDAM Loss 配合 DRW (Deferred Re-Weighting) 训练策略：

python cifar_train.py --gpu 0 --imb_type exp --imb_factor 0.01 --loss_type LDAM --train_rule DRW

参数说明：

• --gpu: 指定使用的 GPU 编号。
• --imb_factor 0.01: 设置不平衡比例为 100 (1/0.01)。
• --loss_type LDAM: 启用标签分布感知边界损失。
• --train_rule DRW: 启用延迟重加权训练规则。

训练过程中，日志和 Tensorboard 数据将自动生成，可用于监控模型收敛情况。