重新思考标签在改善类别不平衡学习中的价值

本仓库包含论文的实现代码：
重新思考标签在改善类别不平衡学习中的价值
杨宇哲 和 许志
第34届神经信息处理系统大会（NeurIPS），2020年
[官网] [arXiv] [论文] [幻灯片] [视频]

如果您觉得这段代码或这一想法有用，请考虑引用我们的工作：

@inproceedings{yang2020rethinking,
  title={Rethinking the Value of Labels for Improving Class-Imbalanced Learning},
  author={Yang, Yuzhe and Xu, Zhi},
  booktitle={Conference on Neural Information Processing Systems (NeurIPS)},
  year={2020}
}

概述

在本工作中，我们从理论和实验两方面证明：无论有标签/无标签数据的不平衡程度如何，以及采用何种基础训练方法，半监督学习（利用未标注数据）和__自监督预训练__（先通过自监督任务对模型进行预训练）都能显著提升在不平衡（长尾）数据集上的性能。

半监督不平衡学习：
利用未标注数据有助于形成更清晰的类别边界，从而实现更好的类别分离，尤其对于尾部类别效果显著。

自监督不平衡学习：
自监督预训练（SSP）有助于缓解测试阶段尾部类别的泄露问题，进而学习到更好的边界和表征。

安装

前置条件

• 下载 CIFAR 和 SVHN 数据集，并将其放置在您的 data_path 中。原始数据将由 imbalance_cifar.py 和 imbalance_svhn.py 转换。
• 下载 ImageNet 和 iNaturalist 2018 数据集，并将其放置在您的 data_path 中。长尾版本将使用 imagenet_inat/data/ 目录下相应子文件夹中的 train/val 切分文件（.txt 文件）创建。
• 根据 ImageNet-LT 和 iNaturalist 2018 的情况，相应地修改 imagenet_inat/main.py 中的 data_root。

依赖项

• PyTorch (>= 1.2，已在 1.4 上测试)
• yaml
• scikit-learn
• TensorboardX

代码概览

主要文件

• train_semi.py：在 CIFAR-LT / SVHN-LT 上使用额外的未标注数据训练模型
• train.py：在 CIFAR-LT / SVHN-LT 上使用（或不使用）SSP 训练模型
• imagenet_inat/main.py：在 ImageNet-LT / iNaturalist 2018 上使用（或不使用）SSP 训练模型
• pretrain_rot.py 和 pretrain_moco.py：分别使用 旋转预测 或 MoCo 进行自监督预训练

主要参数

• --dataset：所选长尾数据集的名称
• --imb_factor：不平衡因子（即论文中不平衡比 \rho 的倒数）
• --imb_factor_unlabel：未标注数据的不平衡因子（即未标注不平衡比 \rho_U 的倒数）
• --pretrained_model：自监督预训练模型的路径
• --resume：恢复检查点的路径（也用于评估）

快速入门

半监督不平衡学习

未标注数据的获取

CIFAR-10-LT：CIFAR-10 的未标注数据是按照 这个仓库 使用 80M TinyImages 准备的。简而言之，首先训练一个数据源模型来区分 CIFAR-10 的各个类别以及“非-CIFAR”类别。然后根据每个类别的预测置信度对图像进行排序，并据此构建未标注（不平衡）数据集。请使用以下链接下载准备好的未标注数据，并将其放置在您的 data_path 中：

• 来自 TinyImages 的 CIFAR-10-LT 未标注数据集

SVHN-LT：由于其自身数据集中包含额外的 531.1K 个样本（已标注），这些样本直接被用来模拟未标注数据集。

需要注意的是，未标注数据中的类别不平衡也被考虑在内，这由 --imb_factor_unlabel 参数控制（即论文中的 \rho_U）。详细信息请参阅 imbalance_cifar.py 和 imbalance_svhn.py。

基于伪标签的半监督学习

要进行伪标签（自训练），首先需要在原始的不平衡数据集上训练一个基础分类器。然后使用该基础分类器生成伪标签：

python gen_pseudolabels.py --resume <ckpt-path> --data_dir <data_path> --output_dir <output_path> --output_filename <save_name>

我们提供了 CIFAR-10-LT 和 SVHN-LT 在 \rho=50 时的伪标签文件，这些文件是基于使用标准交叉熵（CE）损失训练的基础模型生成的：

• CIFAR-10-LT 在 \rho=50 时生成的伪标签
• SVHN-LT 在 \rho=50 时生成的伪标签

例如，在 CIFAR-10-LT 上使用未标注数据进行训练，假设 \rho=50 且 \rho_U=50：

python train_semi.py --dataset cifar10 --imb_factor 0.02 --imb_factor_unlabel 0.02

自监督不平衡学习

自监督预训练（SSP）

要在 CIFAR-10-LT 上进行旋转 SSP，且 \rho=100：

python pretrain_rot.py --dataset cifar10 --imb_factor 0.01

要在 ImageNet-LT 上进行 MoCo SSP：

python pretrain_moco.py --dataset imagenet --data <data_path>

使用 SSP 模型进行网络训练

在 CIFAR-10-LT 上训练，且 \rho=100：

python train.py --dataset cifar10 --imb_factor 0.01 --pretrained_model <path_to_ssp_model>

在 ImageNet-LT / iNaturalist 2018 上训练：

python -m imagenet_inat.main --cfg <path_to_ssp_config> --model_dir <path_to_ssp_model>

结果与模型

所有相关数据和检查点都可以通过 此链接 查看。具体的结果和检查点如下所示。

半监督不平衡学习

CIFAR-10-LT

模型	Top-1 错误率	下载
CE + D_U@5x (\rho=50 和 \rho_U=1)	16.79	ResNet-32
CE + D_U@5x (\rho=50 和 \rho_U=25)	16.88	ResNet-32
CE + D_U@5x (\rho=50 和 \rho_U=50)	18.36	ResNet-32
CE + D_U@5x (\rho=50 和 \rho_U=100)	19.94	ResNet-32

SVHN-LT

模型	Top-1 错误率	下载
CE + D_U@5x (\rho=50 和 \rho_U=1)	13.07	ResNet-32
CE + D_U@5x (\rho=50 和 \rho_U=25)	13.36	ResNet-32
CE + D_U@5x (\rho=50 和 \rho_U=50)	13.16	ResNet-32
CE + D_U@5x (\rho=50 和 \rho_U=100)	14.54	ResNet-32

测试预训练检查点

python train_semi.py --dataset cifar10 --resume <ckpt-path> -e

自监督不平衡学习

CIFAR-10-LT

• 自监督预训练模型（旋转）

  数据集设置	\rho=100	\rho=50	\rho=10
  下载	ResNet-32	ResNet-32	ResNet-32

• 最终模型（200轮）

  模型	\rho	Top-1 错误率	下载
  CE(均匀) + SSP	10	12.28	ResNet-32
  CE(均匀) + SSP	50	21.80	ResNet-32
  CE(均匀) + SSP	100	26.50	ResNet-32
  CE(平衡) + SSP	10	11.57	ResNet-32
  CE(平衡) + SSP	50	19.60	ResNet-32
  CE(平衡) + SSP	100	23.47	ResNet-32

CIFAR-100-LT

• 自监督预训练模型（旋转）

  数据集设置	\rho=100	\rho=50	\rho=10
  下载	ResNet-32	ResNet-32	ResNet-32

• 最终模型（200轮）

  模型	\rho	Top-1 错误率	下载
  CE(均匀) + SSP	10	42.93	ResNet-32
  CE(均匀) + SSP	50	54.96	ResNet-32
  CE(均匀) + SSP	100	59.60	ResNet-32
  CE(平衡) + SSP	10	41.94	ResNet-32
  CE(平衡) + SSP	50	52.91	ResNet-32
  CE(平衡) + SSP	100	56.94	ResNet-32

ImageNet-LT

• 自监督预训练模型（MoCo）
  [ResNet-50]

• 最终模型（90轮）

  模型	Top-1 错误率	下载
  CE(均匀) + SSP	54.4	ResNet-50
  CE(平衡) + SSP	52.4	ResNet-50
  cRT + SSP	48.7	ResNet-50

iNaturalist 2018

• 自监督预训练模型（MoCo）
  [ResNet-50]

• 最终模型（90轮）

  模型	Top-1 错误率	下载
  CE(均匀) + SSP	35.6	ResNet-50
  CE(平衡) + SSP	34.1	ResNet-50
  cRT + SSP	31.9	ResNet-50

测试预训练检查点

# 在CIFAR-10 / CIFAR-100 上测试
python train.py --dataset cifar10 --resume <ckpt-path> -e

# 在ImageNet-LT / iNaturalist 2018 上测试
python -m imagenet_inat.main --cfg <path_to_ssp_config> --model_dir <path_to_model> --test

致谢

本代码部分基于以下开源项目的实现： OpenLongTailRecognition, classifier-balancing, LDAM-DRW, MoCo，以及 semisup-adv。

联系方式

如有任何问题，请随时通过电子邮件（yuzhe@mit.edu & zhixu@mit.edu）或GitHub Issues与我们联系。祝您使用愉快！

imbalanced-semi-self 快速上手指南

本指南旨在帮助开发者快速部署并使用 imbalanced-semi-self 工具，该工具实现了 NeurIPS 2020 论文《Rethinking the Value of Labels for Improving Class-Imbalanced Learning》中的算法，用于解决长尾分布（类别不平衡）下的半监督学习与自监督预训练问题。

环境准备

系统要求

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.6+
• GPU: 支持 CUDA 的 NVIDIA 显卡（可选，但强烈建议用于加速训练）

前置依赖

请确保已安装以下核心库：

• PyTorch: 版本 >= 1.2 (已在 1.4 上测试通过)
• yaml: 用于配置文件解析
• scikit-learn: 用于数据处理与评估
• TensorboardX: 用于可视化训练过程

提示：国内用户建议使用清华源或阿里源加速 PyTorch 及相关包的安装。

pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install pyyaml scikit-learn tensorboardx -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆仓库

   git clone https://github.com/YyzHarry/imbalanced-semi-self.git
   cd imbalanced-semi-self
   

2. 准备数据集 本项目支持 CIFAR, SVHN, ImageNet-LT 和 iNaturalist 2018。需将下载的数据集放置在指定目录，并通过脚本转换为不平衡版本。

   • CIFAR / SVHN: 下载原始数据集后，运行以下脚本生成不平衡版本：

     # 示例：处理 CIFAR 数据集
     python dataset/imbalance_cifar.py
     # 示例：处理 SVHN 数据集
     python dataset/imbalance_svhn.py
     

     注：对于半监督学习所需的无标签数据（如 CIFAR-10-LT），需从提供的 Google Drive 链接 下载预处理好的文件并放入 data_path。

   • ImageNet-LT / iNaturalist 2018: 下载原始数据集，并根据 imagenet_inat/data/ 子文件夹中的 .txt 分割文件构建长尾版本。 注意：需修改 imagenet_inat/main.py 中的 data_root 变量指向你的数据存放路径。

基本使用

本工具主要包含两种核心模式：半监督不平衡学习 和 自监督不平衡学习。以下提供最简化的运行示例。

场景一：半监督不平衡学习 (Semi-Supervised)

利用无标签数据提升长尾数据集（如 CIFAR-10-LT）的分类性能。

1. 生成伪标签 (如果尚未准备好伪标签文件) 使用基础模型为无标签数据生成伪标签：

   python gen_pseudolabels.py --resume <ckpt-path> --data_dir <data_path> --output_dir <output_path> --output_filename <save_name>
   

2. 开始训练 以 CIFAR-10-LT 为例，设置不平衡因子 $\rho=50$ (--imb_factor 0.02) 和无标签数据不平衡因子 $\rho_U=50$ (--imb_factor_unlabel 0.02)：

   python train_semi.py --dataset cifar10 --imb_factor 0.02 --imb_factor_unlabel 0.02
   

场景二：自监督不平衡学习 (Self-Supervised)

先进行自监督预训练（SSP），再在长尾数据上进行微调。

1. 自监督预训练 使用旋转预测（Rotation Prediction）方法在 CIFAR-10-LT ($\rho=100$) 上进行预训练：

   python pretrain_rot.py --dataset cifar10 --imb_factor 0.01
   

   或者在 ImageNet-LT 上使用 MoCo 进行预训练：

   python pretrain_moco.py --dataset imagenet --data <data_path>
   

2. 加载预训练模型进行微调 使用上一步生成的预训练模型权重进行正式训练：

   python train.py --dataset cifar10 --imb_factor 0.01 --pretrained_model <path_to_ssp_model>
   

   针对 ImageNet-LT 或 iNaturalist 2018 的微调命令：

   python -m imagenet_inat.main --cfg <path_to_ssp_config> --model_dir <path_to_ssp_model>
   

参数说明

• --dataset: 选择数据集名称 (如 cifar10, svhn, imagenet)。
• --imb_factor: 不平衡因子（即论文中不平衡比率 $\rho$ 的倒数）。值越小，长尾效应越明显。
• --pretrained_model: 指定自监督预训练模型的路径。
• --resume: 指定断点续训或评估用的检查点路径。