【图像分类】YOLOv5-6.2全新版本：支持图像分类

原创

zstar_ 2022-08-23 16:43:05 博主文章分类：图像分类 ©著作权

文章标签 分类深度学习人工智能数据集图像分类 文章分类 虚拟化云计算

©著作权归作者所有：来自51CTO博客作者zstar_的原创作品，请联系作者获取转载授权，否则将追究法律责任

前言

众所周知，YOLOv5是一款优秀的目标检测模型，但实际上，它也支持图像分类。在6.1版本中，就暗留了classify这个参数，可以在做检测的同时进行分类。

官方仓库地址：https://github.com/ultralytics/yolov5/releases

更新概览

在几天前刚新出的6.2版本中，直接将分类功能单独剥离开来，使其能够直接训练图像分类数据集。

先看看官网公示的更新说明：

分类功能
新增分类功能，并提供各模型在ImageNet上训练过的预训练模型
ClearML日志记录
与开源实验跟踪器ClearML集成。使用pip安装clearml将启用集成，并允许用户跟踪clearml中的每个训练运行。
Deci.ai优化
在Deci上单击一次即可自动编译和量化YOLOv5，从而获得更好的性能
GPU导出基准
可以使用python utils/benchmarks.py --weights yolov5s.pt --device 0来导出Benchmark (mAP and speed)
训练可完全复现
torch>=1.12.0的单GPU YOLOv5训练现在完全可再现，并且可以使用新的–seed参数（默认seed=0）
优化Apple炼丹体验
Apple Metal Performance Shader(MPS：苹果炼丹工具) 支持Apple M1/M2设备

在这些更新中，我最关注的是图像分类功能，那么本篇就来尝试跑通一下。

分类模型效果

下图是官方贴出来的各分类模型对比图，在可以看到在相同的数据集上，YOLOv5x-cls模型取得了最佳的准确率。下列这些模型官方均提供预训练权重。

【图像分类】YOLOv5-6.2全新版本：支持图像分类_深度学习

工程结构

首先看新版本的工程结构，和前几个版本差别不大。主要是多了一个classify文件夹，包含图像分类训练，验证，检测三个函数。

训练结果会保存在runs/train-cls文件夹中。

【图像分类】YOLOv5-6.2全新版本：支持图像分类_数据集_02

数据集下载

在train.py中，提供这段数据集下载程序段：

# Download Dataset
with torch_distributed_zero_first(LOCAL_RANK), WorkingDirectory(ROOT):
    data_dir = data if data.is_dir() else (DATASETS_DIR / data)
    if not data_dir.is_dir():
        LOGGER.info(f'\nDataset not found ⚠️, missing path {data_dir}, attempting download...')
        t = time.time()
        if str(data) == 'imagenet':
            subprocess.run(f"bash {ROOT / 'data/scripts/get_imagenet.sh'}", shell=True, check=True)
        else:
            url = f'https://github.com/ultralytics/yolov5/releases/download/v1.0/{data}.zip'
            download(url, dir=data_dir.parent)
        s = f"Dataset download success ✅ ({time.time() - t:.1f}s), saved to {colorstr('bold', data_dir)}\n"
        LOGGER.info(s)

如果指定的数据集路径下没有数据集，会从官方仓库v1.0的版本中进行下载。这里我建议是手动去进行下载，受限于网络情况，自动下载很容易失败。

可以看到，在官方仓库v1.0版本中的Assets中包含了很多数据集，我下载了最经典的mnist的数据集来做测试。

【图像分类】YOLOv5-6.2全新版本：支持图像分类_深度学习_03

下载完之后，需要在本地进行解压，解压之后的数据集格式如下图所示：

【图像分类】YOLOv5-6.2全新版本：支持图像分类_图像分类_04

禁用wandb

训练之前，可以选择禁用wandb，wandb是和tensorboard类似的数据记录平台，为了防止报错，可以用下面的方式进行禁用。

在终端环境中输入wandb disabled

同时在utils/loggers/wandb/__init__.py中添加wandb = None

try:
    import wandb

    assert hasattr(wandb, '__version__')  # verify package import not local dir
    if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2') and RANK in {0, -1}:
        try:
            wandb_login_success = wandb.login(timeout=30)
        except wandb.errors.UsageError:  # known non-TTY terminal issue
            wandb_login_success = False
        if not wandb_login_success:
            wandb = None
except (ImportError, AssertionError):
    wandb = None
# 添加以下语句
wandb = None

在utils/loggers/wandb/wandb_utils.py中同样添加

try:
    import wandb

    assert hasattr(wandb, '__version__')  # verify package import not local dir
except (ImportError, AssertionError):
    wandb = None
# 添加以下语句
wandb = None