简介

MMPreTrain 是一个基于 PyTorch 的开源预训练工具箱，隶属于 OpenMMLab 项目。

main 分支支持 PyTorch 1.8 及以上版本。

主要特性

• 多种骨干网络和预训练模型
• 丰富的训练策略（监督学习、自监督学习、多模态学习等）
• 训练技巧包
• 大规模训练配置
• 高效且可扩展
• 强大的模型分析与实验工具集
• 多种开箱即用的推理任务。
  • 图像分类
  • 图像字幕生成
  • 视觉问答
  • 视觉定位
  • 检索（图像到图像、文本到图像、图像到文本）

https://github.com/open-mmlab/mmpretrain/assets/26739999/e4dcd3a2-f895-4d1b-a351-fbc74a04e904

最新动态

🌟 v1.2.0 于 2023年04月01日发布

• 支持 LLaVA 1.5。
• 实现了 RAM，并提供了 Gradio 界面。

🌟 v1.1.0 于 2023年12月10日发布

• 支持 Mini-GPT4 训练，并提供中文模型（基于 Baichuan-7B）
• 支持基于 CLIP 的零样本分类。

🌟 v1.0.0 于 2023年04月07日发布

• 支持更多 多模态 算法的推理，例如 LLaVA、MiniGPT-4、Otter 等。
• 支持约 10 种多模态 数据集！
• 新增 iTPN 和 SparK 自监督学习算法。
• 提供 新配置 和 DeepSpeed/FSDP 结合 FlexibleRunner 的示例。以下是相关文档链接：新配置 和 DeepSpeed/FSDP 结合 FlexibleRunner。

🌟 从 MMClassification 升级至 MMPreTrain

• 整合了来自 MMSelfSup 的自监督学习算法，如 MAE、BEiT 等。
• 支持 RIFormer，一种通过移除 token mixer 而实现的简单但有效的视觉骨干网络。
• 重构了数据集流水线可视化。
• 支持 LeViT、XCiT、ViG、ConvNeXt-V2、EVA、RevViT、EfficientnetV2、CLIP、TinyViT 和 MixMIM 等骨干网络。

此次发布引入了一个全新且灵活的训练与测试引擎，但仍处于开发中。欢迎根据 文档 进行尝试。

同时，存在一些破坏性变更。请查阅 迁移教程。

更多详情及其他版本历史，请参阅 changelog。

安装

以下是快速安装步骤：

conda create -n open-mmlab python=3.8 pytorch==1.10.1 torchvision==0.11.2 cudatoolkit=11.3 -c pytorch -y
conda activate open-mmlab
pip install openmim
git clone https://github.com/open-mmlab/mmpretrain.git
cd mmpretrain
mim install -e .

更多详细的安装和数据集准备信息，请参阅安装文档。

若需支持多模态模型，请通过以下命令安装额外依赖：

mim install -e ".[multimodal]"

用户指南

我们为新用户提供了一系列关于 MMPreTrain 基本使用的教程：

• 配置文件简介
• 数据集准备
• 使用现有模型进行推理
• 训练
• 测试
• 下游任务

如需更多信息，请参阅我们的文档。

模型动物园

结果和模型可在模型动物园中找到。

支持的骨干网络	自监督学习	多模态算法	其他
VGG ResNet ResNeXt SE-ResNet SE-ResNeXt RegNet ShuffleNet V1 ShuffleNet V2 MobileNet V2 MobileNet V3 Swin-Transformer Swin-Transformer V2 RepVGG Vision-Transformer Transformer-in-Transformer Res2Net MLP-Mixer DeiT DeiT-3 Conformer T2T-ViT Twins EfficientNet EdgeNeXt ConvNeXt HRNet VAN ConvMixer CSPNet PoolFormer Inception V3 MobileOne EfficientFormer MViT HorNet MobileViT DaViT RepLKNet BEiT MixMIM EfficientNet V2 RevViT ConvNeXt V2 ViG XCiT LeViT RIFormer GLIP ViT SAM EVA02 DINO V2 HiViT	MoCo V1 (CVPR'2020) SimCLR (ICML'2020) MoCo V2 (arXiv'2020) BYOL (NeurIPS'2020) SwAV (NeurIPS'2020) DenseCL (CVPR'2021) SimSiam (CVPR'2021) Barlow Twins (ICML'2021) MoCo V3 (ICCV'2021) BEiT (ICLR'2022) MAE (CVPR'2022) SimMIM (CVPR'2022) MaskFeat (CVPR'2022) CAE (arXiv'2022) MILAN (arXiv'2022) BEiT V2 (arXiv'2022) EVA (CVPR'2023) MixMIM (arXiv'2022) iTPN (CVPR'2023) SparK (ICLR'2023) MFF (ICCV'2023)	BLIP (arxiv'2022) BLIP-2 (arxiv'2023) OFA (CoRR'2022) Flamingo (NeurIPS'2022) Chinese CLIP (arxiv'2022) MiniGPT-4 (arxiv'2023) LLaVA (arxiv'2023) Otter (arxiv'2023)	图像检索任务： ArcFace (CVPR'2019) 训练与测试技巧： RandAug AutoAug RepeatAugSampler TTA ...

贡献

我们感谢所有为改进 MMPreTrain 所做的贡献。请参阅CONTRIBUTING以获取贡献指南。

致谢

MMPreTrain 是一个开源项目，由来自不同院校和公司的研究人员及工程师共同贡献而成。我们感谢所有实现其方法或添加新功能的贡献者，以及提供宝贵反馈的用户。我们希望这个工具箱和基准能够服务于不断壮大的研究社区，为重现实验室现有方法并支持他们自身的学术研究提供灵活的工具包。

引用

如果您在研究中使用了本项目，请考虑引用：

@misc{2023mmpretrain,
    title={OpenMMLab预训练工具箱与基准},
    author={MMPreTrain贡献者},
    howpublished = {\url{https://github.com/open-mmlab/mmpretrain}},
    year={2023}
}

许可证

本项目采用Apache 2.0许可证发布。

OpenMMLab中的项目

• MMEngine：OpenMMLab用于训练深度学习模型的基础库。
• MMCV：OpenMMLab用于计算机视觉的基础库。
• MIM：MIM用于安装OpenMMLab的相关软件包。
• MMEval：一个面向多种机器学习框架的统一评估库。
• MMPreTrain：OpenMMLab预训练工具箱及基准。
• MMDetection：OpenMMLab目标检测工具箱及基准。
• MMDetection3D：OpenMMLab新一代通用3D目标检测平台。
• MMRotate：OpenMMLab旋转目标检测工具箱及基准。
• MMyolo：OpenMMLab YOLO系列工具箱及基准。
• MMSegmentation：OpenMMLab语义分割工具箱及基准。
• MMOCR：OpenMMLab文本检测、识别与理解工具箱。
• MMPose：OpenMMLab人体姿态估计工具箱及基准。
• MMHuman3D：OpenMMLab 3D人体参数化模型工具箱及基准。
• MMSelfSup：OpenMMLab自监督学习工具箱及基准。
• MMRazor：OpenMMLab模型压缩工具箱及基准。
• MMFewShot：OpenMMLab小样本学习工具箱及基准。
• MMAction2：OpenMMLab新一代动作理解工具箱及基准。
• MMTracking：OpenMMLab视频感知工具箱及基准。
• MMFlow：OpenMMLab光流计算工具箱及基准。
• MMagic：OpenMMLab Advanced, Generative and Intelligent Creation工具箱。
• MMGeneration：OpenMMLab图像和视频生成模型工具箱。
• MMDeploy：OpenMMLab模型部署框架。
• Playground：一个汇集并展示基于OpenMMLab构建的优秀项目的中心平台。

conda create -n open-mmlab python=3.8 pytorch==1.10.1 torchvision==0.11.2 cudatoolkit=11.3 -c pytorch -y
conda activate open-mmlab

pip install openmim

git clone https://github.com/open-mmlab/mmpretrain.git
cd mmpretrain
mim install -e .

mim install -e ".[multimodal]"

from mmpretrain import inference_model

# 指定模型配置名称和图片路径
# 这里以 ResNet50 在 ImageNet 上预训练的模型为例
result = inference_model('resnet50_8xb32_in1k', 'demo.jpg')

# 打印预测结果
print(result['pred_class'])  # 输出预测类别名称
print(result['pred_score'])  # 输出置信度分数