爱笑的小姐姐 · 2022年08月19日 · 北京市

Yolov5 也可以训练分类模型啦,语义分割+实例分割很快到来

image.png

yolov5-6.2增加了分类训练、验证、预测和导出(所有 11 种格式),还提供了 ImageNet 预训练的 YOLOv5m-cls、ResNet(18、34、50、101) 和 EfficientNet (b0-b3) 模型。
 
此次发布的主要目标是引入超级简单的 YOLOv5 分类工作流程,就像现有的目标检测模型一样。以下新的 v6.2 YOLOv5-cls 模型只是一个开始,作者将继续与现有的检测模型一起改进这些模型。
 
下一个版本 v6.3 计划于 9 月发布,将为 YOLOv5 带来官方实例分割支持,今年晚些时候将发布一个主要的 v7.0 版本,更新所有 3 个任务的架构——分类、检测和分割。

1、重要更新

分类模型:TensorFlow、Keras、TFLite、TF.js 模型导出现在使用 python export.py --include saved_model pb tflite tfjs 完全集成。

ClearML日志记录:与开源实验跟踪器 ClearML 集成。使用 pip install clearml 安装将启用集成并允许用户跟踪在 ClearML 中运行的每个训练。这反过来又允许用户跟踪和比较运行,甚至远程安排运行。

Deci.ai优化:一键自动编译和量化 YOLOv5 以获得更好的推理性能。

GPU导出基准:使用 python utils/benchmarks.py --weights yolov5s.pt --device 0 用于 GPU 基准测试或 --device cpu 用于 CPU 基准测试,对所有 YOLOv5 导出格式进行基准测试(mAP 和速度)。

训练可再现性:使用 torch>=1.12.0 的单 GPU YOLOv5 训练现在完全可再现,并且可以使用新的 --seed 参数(默认种子 = 0)。

Apple Metal Performance Shader (MPS) 支持:通过 --device mps 对 Apple M1/M2 设备的 MPS 支持(完整功能在 pytorch/pytorch#77764 中等待更新)。

2、分类模型与精度

使用 4×A100 在 ImageNet 上训练了 YOLOv5-cls 分类模型 90 个 epoch,并且训练了 ResNet 和 EfficientNet 模型以及相同的默认训练设置进行比较。将所有模型导出到 ONNX FP32 进行 CPU 速度测试,并将所有模型导出到 TensorRT FP16 进行 GPU 速度测试。在 Google Colab Pro 上进行了所有速度测试,以便轻松重现。

image.png

3、使用

YOLOv5 分类训练支持使用 --data 参数自动下载 MNIST、Fashion-MNIST、CIFAR10、CIFAR100、Imagenette、Imagewoof 和 ImageNet 数据集。例如,要开始在 MNIST 上进行训练,使用 --data mnist。

train

`# Single-GPU  
python classify/train.py --model yolov5s-cls.pt --data cifar100 --epochs 5 --img 224 --batch 128  
  
# Multi-GPU DDP  
python -m torch.distributed.run --nproc_per_node 4 --master_port 1 classify/train.py --model yolov5s-cls.pt --data imagenet --epochs 5 --img 224 --device 0,1,2,3  
`

val


`bash data/scripts/get_imagenet.sh --val  # download ImageNet val split (6.3G, 50000 images)  
python classify/val.py --weights yolov5m-cls.pt --data ../datasets/imagenet --img 224  # validate  
`

test

`python classify/predict.py --weights yolov5s-cls.pt --data data/images/bus.jpg  
`

4、构建形式

分类模型的构建依旧是YOLOv5的风格,加入了分类的head,这里点赞,不怕没有预训练权重了!

`class ClassificationModel(BaseModel):  
    # YOLOv5 classification model  
    def __init__(self, cfg=None, model=None, nc=1000, cutoff=10):  # yaml, model, number of classes, cutoff index  
        super().__init__()  
        self._from_detection_model(model, nc, cutoff) if model is not None else self._from_yaml(cfg)  
  
    def _from_detection_model(self, model, nc=1000, cutoff=10):  
        # Create a YOLOv5 classification model from a YOLOv5 detection model  
        if isinstance(model, DetectMultiBackend):  
            model = model.model  # unwrap DetectMultiBackend  
        model.model = model.model[:cutoff]  # backbone  
        m = model.model[-1]  # last layer  
        ch = m.conv.in_channels if hasattr(m, 'conv') else m.cv1.conv.in_channels  # ch into module  
        c = Classify(ch, nc)  # Classify()  
        c.i, c.f, c.type = m.i, m.f, 'models.common.Classify'  # index, from, type  
        model.model[-1] = c  # replace  
        self.model = model.model  
        self.stride = model.stride  
        self.save = []  
        self.nc = nc  
  
    def _from_yaml(self, cfg):  
        # Create a YOLOv5 classification model from a *.yaml file  
        self.model = None  
`

参考

https://github.com/ultralytics/yolov5/releases

作者:ChaucerG
文章来源:集智书童

推荐阅读

更多汽车电子干货请关注 嵌入式AI 专栏。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
18801
内容数
1347
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息