MindCV: 基于MindSpore的计算机视觉工具箱

MindCV简介

MindCV是一个基于MindSpore深度学习框架的开源计算机视觉工具箱。它由MindSpore团队、西安电子科技大学和西安交通大学联合开发,旨在为计算机视觉研究和应用提供丰富的资源和便捷的开发环境。

MindCV具有以下主要特点:

易用性强。MindCV将视觉框架分解为各种可配置的组件,用户可以轻松自定义数据管道、模型和学习流程:

import mindcv
# 创建数据集
dataset = mindcv.create_dataset('cifar10', download=True)
# 创建模型
network = mindcv.create_model('resnet50', pretrained=True)

用户还可以通过一行命令定制和启动迁移学习或训练任务:

python train.py --model=swin_tiny --pretrained --opt=adamw --lr=0.001 --data_dir=/path/to/data

包含最先进模型。MindCV提供了各种基于CNN和Transformer的视觉模型,包括SwinTransformer。同时提供预训练权重和性能报告,帮助用户选择和复用合适的模型。
灵活高效。MindCV基于MindSpore构建,可以在不同硬件平台(GPU/CPU/Ascend)上运行。它同时支持高效的图模式和灵活的动态图模式。

模型库

MindCV提供了丰富的预训练模型,包括:

ResNet系列
MobileNet系列
VGG
Inception系列
DenseNet
EfficientNet系列
Swin Transformer
Vision Transformer (ViT)
等多种经典和最新的CV模型

这些模型的性能指标、训练配方和预训练权重都可以在模型库页面查看。

Image 1: a yellow labrador standing on the grass

快速上手

安装

MindCV的安装非常简单,详细步骤请参考安装指南。

使用示例

下面是几个简单的代码片段,展示了MindCV的基本用法:

import mindcv

# 列出并查找预训练的视觉模型
mindcv.list_models("swin*", pretrained=True)
# ['swin_tiny']

# 创建模型对象
network = mindcv.create_model('swin_tiny', pretrained=True)

# 验证模型精度
!python validate.py --model=swin_tiny --pretrained --dataset=imagenet --val_split=validation
# {'Top_1_Accuracy': 0.80824, 'Top_5_Accuracy': 0.94802, 'loss': 1.7331367141008378}

图像分类演示

以下是使用预训练SoTA模型对图像进行分类的示例:

python infer.py --model=swin_tiny --image_path='./dog.jpg'
# {'Labrador retriever': 0.5700152, 'golden retriever': 0.034551315, 'kelpie': 0.010108651, 'Chesapeake Bay retriever': 0.008229004, 'Walker hound, Walker foxhound': 0.007791956}

top-1预测结果是拉布拉多猎犬,这与图中的品种相符。

训练

MindCV提供了灵活的训练脚本train.py,可以轻松地在标准或自定义数据集上训练模型。训练策略(如数据增强、学习率调度等)可以通过外部参数或yaml配置文件进行配置。

单机训练

python train.py --model=resnet50 --dataset=cifar10 --dataset_download

分布式训练

对于ImageNet等大型数据集,可以使用mpirun在多个设备上进行分布式训练:

mpirun -n 4 python train.py --distribute \
    --model=densenet121 --dataset=imagenet --data_dir=/path/to/imagenet

配置和训练策略

MindCV提供了20多个在ImageNet上实现SoTA结果的训练配方。用户可以通过修改yaml文件轻松地将这些训练策略应用到自己的模型上以提高性能。

验证

使用validate.py评估模型性能:

python validate.py --model=resnet50 --dataset=imagenet --data_dir=/path/to/data --ckpt_path=/path/to/model.ckpt

还可以通过启用--val_while_train选项在训练过程中跟踪验证精度。

支持的算法

MindCV支持多种先进的计算机视觉算法,包括:

数据增强:AutoAugment、RandAugment、Repeated Augmentation等
优化器:Adam、AdamW、Lion等
学习率调度器:Warmup Cosine Decay、Step LR等
正则化:Weight Decay、Label Smoothing等
损失函数:Cross Entropy、Binary Cross Entropy等
集成学习:Warmup EMA

贡献指南

我们欢迎各种形式的贡献,包括问题报告和PR,以使MindCV变得更好。请参考贡献指南了解详细信息。

许可证

本项目遵循Apache License 2.0开源许可证。

MindCV是一个功能强大、易用性高的计算机视觉工具箱,为研究人员和开发者提供了丰富的资源和便捷的开发环境。无论是进行学术研究还是工业应用,MindCV都是一个值得尝试的优秀工具。