ImageNet21K预训练模型使用指南
项目地址: https://gitcode.com/gh_mirrors/im/ImageNet21K 项目概述
ImageNet21K是阿里巴巴达摩院MIIL团队开发的基于ImageNet-21K数据集的大规模预训练模型项目。该项目基于论文《ImageNet-21K Pretraining for the Masses》（NeurIPS 2021），通过专门的预处理流程、WordNet层次结构利用以及新颖的语义softmax训练方案，使得各种模型（包括小型移动端模型）在多个数据集和任务上都能显著受益于ImageNet-21K预训练。
项目目录结构
.
├── README.md # 项目说明文档
├── LICENSE # 许可证文件
├── MODEL_ZOO.md # 预训练模型库说明
├── Transfer_learning.md # 迁移学习说明文档
├── requirements.txt # Python依赖包列表
├── dataset_preprocessing/ # 数据集预处理相关文件
│ ├── processing_instructions.md
│ ├── processing_script.sh
│ ├── processing_script_short_edge.sh
│ ├── resize.py
│ ├── resize_short_edge.py
│ └── validating_files.py
├── pics/ # 项目图片资源
│ ├── dog_inference.png
│ ├── main_pic.png
│ ├── ms_coco_scores.png
│ ├── pic1.png
│ ├── pic2.png
│ └── pic3.png
├── src_files/ # 核心源码目录
│ ├── data_loading/ # 数据加载模块
│ │ └── data_loader.py
│ ├── helper_functions/ # 辅助功能模块
│ │ ├── augmentations.py
│ │ ├── distributed.py
│ │ └── general_helper_functions.py
│ ├── loss_functions/ # 损失函数模块
│ │ └── losses.py
│ ├── models/ # 模型定义目录
│ │ ├── __init__.py
│ │ ├── ofa/ # OFA模型实现
│ │ ├── tresnet/ # TResNet模型实现
│ │ └── utils/ # 模型工具模块
│ ├── optimizers/ # 优化器模块
│ │ └── create_optimizer.py
│ └── semantic/ # 语义处理模块
│ ├── metrics.py
│ ├── semantic_loss.py
│ └── semantics.py
├── tests/ # 测试目录
│ └── test_semantic_softmax_loss.py
├── train_semantic_softmax.py # 语义softmax训练脚本
├── train_single_label.py # 单标签训练脚本
├── train_single_label_from_scratch.py # 从零开始训练脚本
└── visualize_detector.py # 可视化推理脚本
预训练模型库
项目提供了多个在ImageNet-21K-P数据集上使用语义softmax训练的预训练模型：
| 模型 | ImageNet-21K-P语义准确率 | ImageNet-1K准确率 | 最大批次大小 | 训练速度(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 |
快速开始
1. 安装依赖
首先安装必要的Python依赖包：
pip install -r requirements.txt
2. 使用预训练模型
通过timm库可以直接使用预训练模型：
import timm
# ImageNet-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)
# ImageNet-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)
3. 模型训练
使用语义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
4. 可视化推理
使用可视化脚本进行推理演示：
python visualize_detector.py
数据集处理
数据集处理说明和脚本位于dataset_preprocessing目录下，包括图像大小调整、验证文件检查等功能：
- processing_instructions.md：详细的数据集处理说明
- processing_script.sh：处理脚本
- resize.py：图像大小调整脚本
- validating_files.py：文件验证脚本
迁移学习优势
项目在多个数据集上取得了State-of-the-Art的结果，包括MS-COCO、Pascal-VOC、Stanford Cars和CIFAR-100。更重要的是，ImageNet-21K预训练使得模型在迁移学习中更加稳定，对超参数选择不那么敏感。
核心特性
- 语义softmax训练：利用WordNet层次结构信息，提高模型性能
- 高效预处理：专门的预处理流程使ImageNet-21K预训练对大众可用
- 多模型支持：支持MobilenetV3、TResNet、ViT、OFA等多种模型架构
- 迁移学习优化：显著提升下游任务的迁移学习效果和稳定性
技术实现
项目采用PyTorch实现，主要技术特点包括：
- 分布式训练支持
- 混合精度训练加速
- 语义层次结构利用
- 多种优化器支持
- 完整的训练和评估流程
通过该项目，研究人员和开发者可以轻松利用ImageNet-21K的大规模数据进行预训练，从而在各种计算机视觉任务中获得更好的性能表现。
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
