面向大众的ImageNet-21K预训练

论文 | 预训练模型

官方 PyTorch 实现

Tal Ridnik, Emanuel Ben-Baruch, Asaf Noy, Lihi Zelnik-Manor
达摩院，阿里巴巴集团

摘要

ImageNet-1K 是用于计算机视觉任务深度学习模型预训练的主要数据集。而包含更多图片和类别的 ImageNet-21K 数据集，由于其复杂性以及人们对其相较于标准 ImageNet-1K 预训练所带来价值的低估，较少被用于预训练。

本文旨在填补这一空白，使高质量、高效的 ImageNet-21K 预训练能够惠及所有人。通过专门的预处理阶段、利用 WordNet 层次结构，以及一种称为语义 softmax 的新颖训练方案，我们证明了包括小型移动端模型在内的多种模型，在众多数据集和任务上都能显著受益于 ImageNet-21K 预训练。

此外，我们还表明，对于 ViT 等新兴的前沿模型，我们的方法优于以往的 ImageNet-21K 预训练方案。我们提出的预训练流程高效且易于使用，并能基于公开可用的数据集获得可复现的最先进结果。

2023年10月1日更新

新增由 CodiumAI 自动生成的测试用例 tests

2021年11月29日更新 - 新文章发布，提供具有最先进性能的新分类头

请查看我们的新项目 Ml-Decoder，它为多标签、单标签及零样本任务提供统一的分类头。搭载 ML-Decoder 的骨干网络不仅达到了 SOTA 水平，还改善了速度与精度之间的权衡。

2021年8月1日更新 - “面向大众的 ImageNet-21K 预训练”已被 NeurIPS 2021 接受！

我们非常高兴地宣布，“面向大众的 ImageNet-21K 预训练”一文已被 NeurIPS 2021 接受（数据集与基准测试 赛道）。OpenReview 页面可在 这里 查看。

入门指南

(0) 可视化与推理脚本

首先，您可以使用以下 script 对特定图像进行玩耍和推理。示例如下：

(1) ImageNet-21K-P 数据集上的预训练模型

主干网络	ImageNet-21K-P 语义 top-1 准确率 [%]	ImageNet-1K top-1 准确率 [%]	最大 批量大小	最大 训练速度 (img/sec)	最大 推理速度 (img/sec)
MobilenetV3_large_100	73.1	78.0	488	1210	5980
OFA_flops_595m_s	75.0	81.0	288	500	3240
ResNet50	75.6	82.0	320	720	2760
TResNet-M	76.4	83.1	520	670	2970
TResNet-L (V2)	76.7	83.9	240	300	1460
ViT-B-16	77.6	84.4	160	340	1140

更多详情请参见 这里。
我们强烈建议从测试这些权重与标准 ImageNet-1K 预训练的对比开始，并在您相关的下游任务上比较结果。当您看到显著提升后，再继续进行新模型的预训练。

请注意，我们的一些同时具备 21K 和 1K 预训练的模型也可以通过优秀的 timm 包获得：

21K:
model = timm.create_model('mobilenetv3_large_100_miil_in21k', pretrained=True)
model = timm.create_model('tresnet_m_miil_in21k', pretrained=True)
model = timm.create_model('vit_base_patch16_224_miil_in21k', pretrained=True)
model = timm.create_model('mixer_b16_224_miil_in21k', pretrained=True)


1K:
model = timm.create_model('mobilenetv3_large_100_miil', pretrained=True)
model = timm.create_model('tresnet_m', pretrained=True)
model = timm.create_model('vit_base_patch16_224_miil', pretrained=True)
model = timm.create_model('mixer_b16_224_miil', pretrained=True)

您可以使用此 链接 来确认我们确实达到了文章中报告的准确率。

(2) 数据集的获取与处理

有关数据集的获取和处理说明，请参阅 这里。

(3) 训练代码

要使用训练代码，首先将相关的 ImageNet-21K-P 语义树文件下载到本地 ./resources/ 文件夹中：

• fall11 语义树
• winter21 语义树

语义 softmax 训练示例：

python train_semantic_softmax.py \
--batch_size=4 \
--data_path=/mnt/datasets/21k \
--model_name=mobilenetv3_large_100 \
--model_path=/mnt/models/mobilenetv3_large_100.pth \
--epochs=80

为了缩短训练时间，我们从标准 ImageNet-1K 的权重进行初始化。建议使用 timm 仓库 中的 ImageNet-1K 权重。

(4) 迁移学习代码

有关复现代码，请参阅 这里，其中展示了 miil 预训练不仅能够提升迁移学习的效果，还能使 MLP 模型在超参数选择方面更加稳定和鲁棒。

其他 SoTA 结果

文章中的结果是在固定超参数下的对比结果。此外，借助我们的预训练模型以及针对不同数据集（分辨率、优化器、学习率）调整超参数的专用训练方案，我们在多个计算机视觉数据集上——MS-COCO、Pascal-VOC、Stanford Cars 和 CIFAR-100——都取得了 SoTA 水平的结果。

我们将分享模型的检查点以验证我们的成绩。

引用

@misc{ridnik2021imagenet21k,
      title={ImageNet-21K Pretraining for the Masses}, 
      author={Tal Ridnik and Emanuel Ben-Baruch and Asaf Noy and Lihi Zelnik-Manor},
      year={2021},
      eprint={2104.10972},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

ImageNet-21K 快速上手指南

ImageNet-21K 是一个包含 21,000 个类别的大规模图像数据集，旨在为计算机视觉任务提供高质量的预训练模型。本项目提供了基于 PyTorch 的官方实现，通过语义 Softmax（Semantic Softmax）等创新技术，显著提升了模型在多种下游任务中的表现。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu 18.04+)
• Python: 3.7 或更高版本
• GPU: 支持 CUDA 的 NVIDIA 显卡（建议显存 16GB+ 以运行大模型）

前置依赖

确保已安装以下核心库：

• PyTorch (1.7+)
• torchvision
• timm (用于加载预训练模型)

安装基础依赖命令：

pip install torch torchvision timm

安装步骤

本项目主要作为代码库和权重资源使用，无需复杂的编译安装。只需克隆仓库并准备数据即可。

1. 克隆项目代码

git clone https://github.com/Alibaba-MIIL/ImageNet21K.git
cd ImageNet21K

2. 下载预训练权重（可选） 如果仅需使用预训练模型进行推理或迁移学习，可直接通过 timm 库自动下载，或手动从提供的链接下载权重文件至本地。

基本使用

方式一：直接使用预训练模型（推荐）

最快捷的方式是通过 timm 库直接加载已在 ImageNet-21K 上预训练好的模型。支持 MobileNetV3, ResNet50, ViT 等主流骨干网络。

import timm

# 加载 ImageNet-21K 预训练模型 (示例：MobileNetV3)
model = timm.create_model('mobilenetv3_large_100_miil_in21k', pretrained=True)

# 加载 ImageNet-1K 预训练模型 (对比用)
# model = timm.create_model('mobilenetv3_large_100_miil', pretrained=True)

model.eval()
# 接下来可进行推理或微调

支持的模型列表包括：

• mobilenetv3_large_100_miil_in21k
• tresnet_m_miil_in21k
• vit_base_patch16_224_miil_in21k
• mixer_b16_224_miil_in21k

方式二：运行推理与可视化

使用项目自带的脚本对图片进行推理测试：

python visualize_detector.py

方式三：从头开始训练（语义 Softmax）

若需使用 ImageNet-21K 数据集重新训练模型，需先完成以下步骤：

1. 准备数据集 参考 dataset_preprocessing/processing_instructions.md 获取并处理 ImageNet-21K-P 数据集。

2. 下载语义树文件 将对应的语义树文件下载到本地 ./resources/ 文件夹：

   • fall11 版本: imagenet21k_miil_tree.pth
   • winter21 版本: imagenet21k_miil_tree.pth

3. 执行训练 以下示例展示如何使用语义 Softmax 进行训练（基于 ImageNet-1K 权重初始化以加速收敛）：

python train_semantic_softmax.py \
--batch_size=4 \
--data_path=/mnt/datasets/21k \
--model_name=mobilenetv3_large_100 \
--model_path=/mnt/models/mobilenetv3_large_100.pth \
--epochs=80

提示：建议先在您的下游任务上对比测试预训练权重与标准 ImageNet-1K 权重的效果，确认提升后再进行大规模预训练。