一、训练yolox模型

  • 下载yolox
git clone https://github.com/Megvii-BaseDetection/YOLOX.git

下载yolox-nano的权重文件:https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_nano.pth

  • 环境

创建一个conda虚拟环境yolox,主要安装YOLOX/requirements.txt下载必要的库

conda activate yolox
cd YOLOX
python setup.py develop
  • 配置文件
  1. 在路径YOLOX/exps/example/yolox_voc下创建一个文件yolox_voc_nano.py
# encoding: utf-8
import os

import torch
import torch.distributed as dist

import torch.nn as nn

from yolox.data import get_yolox_datadir
from yolox.exp import Exp as MyExp


class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.num_classes = 3 # 类别数量

        self.depth = 0.33
        self.width = 0.25
        self.input_size = (416, 416)
        self.random_size = (10, 20)
        self.mosaic_scale = (0.5, 1.5)
        self.test_size = (416, 416)
        self.max_epoch = 500
        self.enable_mixup = False
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

    def get_model(self, sublinear=False):

        def init_yolo(M):
            for m in M.modules():
                if isinstance(m, nn.BatchNorm2d):
                    m.eps = 1e-3
                    m.momentum = 0.03
        if "model" not in self.__dict__:
            from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
            in_channels = [256, 512, 1024]
            # NANO model use depthwise = True, which is main difference.
            backbone = YOLOPAFPN(
                self.depth, self.width, in_channels=in_channels,
                act=self.act, depthwise=True,
            )
            head = YOLOXHead(
                self.num_classes, self.width, in_channels=in_channels,
                act=self.act, depthwise=True
            )
            self.model = YOLOX(backbone, head)

        self.model.apply(init_yolo)
        self.model.head.initialize_biases(1e-2)
        return self.model

    def get_data_loader(self, batch_size, is_distributed, no_aug=False, cache_img=False):
        from yolox.data import (
            VOCDetection,
            TrainTransform,
            YoloBatchSampler,
            DataLoader,
            InfiniteSampler,
            MosaicDetection,
            worker_init_reset_seed,
        )
        from yolox.utils import (
            wait_for_the_master,
            get_local_rank,
        )
        local_rank = get_local_rank()

        with wait_for_the_master(local_rank):
            dataset = VOCDetection(
                data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
                image_sets=[('2007', 'trainval')],
                img_size=self.input_size,
                preproc=TrainTransform(
                    max_labels=50,
                    flip_prob=self.flip_prob,
                    hsv_prob=self.hsv_prob),
                cache=cache_img,
            )

        dataset = MosaicDetection(
            dataset,
            mosaic=not no_aug,
            img_size=self.input_size,
            preproc=TrainTransform(
                max_labels=120,
                flip_prob=self.flip_prob,
                hsv_prob=self.hsv_prob),
            degrees=self.degrees,
            translate=self.translate,
            mosaic_scale=self.mosaic_scale,
            mixup_scale=self.mixup_scale,
            shear=self.shear,
            enable_mixup=self.enable_mixup,
            mosaic_prob=self.mosaic_prob,
            mixup_prob=self.mixup_prob,
        )

        self.dataset = dataset

        if is_distributed:
            batch_size = batch_size // dist.get_world_size()

        sampler = InfiniteSampler(
            len(self.dataset), seed=self.seed if self.seed else 0
        )

        batch_sampler = YoloBatchSampler(
            sampler=sampler,
            batch_size=batch_size,
            drop_last=False,
            mosaic=not no_aug,
        )

        dataloader_kwargs = {"num_workers": self.data_num_workers, "pin_memory": True}
        dataloader_kwargs["batch_sampler"] = batch_sampler

        # Make sure each process has different random seed, especially for 'fork' method
        dataloader_kwargs["worker_init_fn"] = worker_init_reset_seed

        train_loader = DataLoader(self.dataset, **dataloader_kwargs)

        return train_loader

    def get_eval_loader(self, batch_size, is_distributed, testdev=False, legacy=False):
        from yolox.data import VOCDetection, ValTransform

        valdataset = VOCDetection(
            data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
            image_sets=[('2007', 'test')],
            img_size=self.test_size,
            preproc=ValTransform(legacy=legacy),
        )

        if is_distributed:
            batch_size = batch_size // dist.get_world_size()
            sampler = torch.utils.data.distributed.DistributedSampler(
                valdataset, shuffle=False
            )
        else:
            sampler = torch.utils.data.SequentialSampler(valdataset)

        dataloader_kwargs = {
            "num_workers": self.data_num_workers,
            "pin_memory": True,
            "sampler": sampler,
        }
        dataloader_kwargs["batch_size"] = batch_size
        val_loader = torch.utils.data.DataLoader(valdataset, **dataloader_kwargs)

        return val_loader

    def get_evaluator(self, batch_size, is_distributed, testdev=False, legacy=False):
        from yolox.evaluators import VOCEvaluator

        val_loader = self.get_eval_loader(batch_size, is_distributed, testdev, legacy)
        evaluator = VOCEvaluator(
            dataloader=val_loader,
            img_size=self.test_size,
            confthre=self.test_conf,
            nmsthre=self.nmsthre,
            num_classes=self.num_classes,
        )
        return evaluator

  1. 修改YOLOX/yolox/data/datasets/voc_classes.py中VOC_CLASSES为自己的类别名称。

  2. 如果训练电脑为windows系统的还需要修改YOLOX/yolox/data/datasets/voc.py

# 第一处原始
path_filename = [
            (self._imgpath % self.ids[i]).split(self.root + "/")[1]
            for i in range(self.num_imgs)
        ]
# 第一处修改
    path_filename = [
            (self._imgpath % self.ids[i]).split(self.root + "\\")[1]
            for i in range(self.num_imgs)
        ]
# 第二处原始
annopath = os.path.join(rootpath, "Annotations", "{:s}.xml")
# 第二处修改
annopath = os.path.join(rootpath, "Annotations") + "\\{:s}.xml"

否则训练的时候会报错如下

  # 第一处报错
  File "e:\projects\codes\yolox\yolox\data\datasets\voc.py", line 137, in __init__
    path_filename = [

  File "e:\projects\codes\yolox\yolox\data\datasets\voc.py", line 138, in <listcomp>
    (self._imgpath % self.ids[i]).split(self.root + "/")[1]
     │    │          │    │   │         │    └ 'e:\\projects\\codes\\yolox\\datasets\\VOCdevkit'
     │    │          │    │   │         └ <yolox.data.datasets.voc.VOCDetection object at 0x0000021DD1A8C610>
     │    │          │    │   └ 0
     │    │          │    └ [('e:\\projects\\codes\\yolox\\datasets\\VOCdevkit\\VOC2007', 'bc-057'), ('e:\\projects\\codes\\yolox\\...
     │    │          └ <yolox.data.datasets.voc.VOCDetection object at 0x0000021DD1A8C610>
     │    └ '%s\\JPEGImages\\%s.jpg'
     └ <yolox.data.datasets.voc.VOCDetection object at 0x0000021DD1A8C610>

IndexError: list index out of range
# 第二处报错
  File "e:\projects\codes\yolox\yolox\evaluators\voc_eval.py", line 16, in parse_rec
    tree = ET.parse(filename)
           │  │     └ 'bc-064.xml'
           │  └ <function parse at 0x0000022CFD6D15E0>
           └ <module 'xml.etree.ElementTree' from 'F:\\Anaconda3\\envs\\yolox\\lib\\xml\\etree\\ElementTree.py'>

  File "F:\Anaconda3\envs\yolox\lib\xml\etree\ElementTree.py", line 1202, in parse
    tree.parse(source, parser)
    │    │     │       └ None
    │    │     └ 'bc-064.xml'
    │    └ <function ElementTree.parse at 0x0000022CFD6D0670>
    └ <xml.etree.ElementTree.ElementTree object at 0x0000022CFFDB20A0>

  File "F:\Anaconda3\envs\yolox\lib\xml\etree\ElementTree.py", line 584, in parse
    source = open(source, "rb")
                  └ 'bc-064.xml'

FileNotFoundError: [Errno 2] No such file or directory: 'bc-064.xml'
  1. 如果数据集中存在不同格式的图片,则还需要修改下面部分YOLOX/yolox/data/datasets/voc.py
# 第一处原始
		self._imgpath = os.path.join("%s", "JPEGImages", "%s.jpg")
# 第一处修改
        self._imgpath = os.path.join("%s", "JPEGImages", "%s")
# 第二处原始
        path_filename = [
            (self._imgpath % self.ids[i]).split(self.root + "\\")[1]
            for i in range(self.num_imgs)
        ]
# 第二处修改
        path_filename = []
        img_names = ['.jpg','.png','.bmp', '.jpeg']
        for i in range(self.num_imgs):
            for n in img_names:
                if os.path.exists(self._imgpath % (self.ids[i][0], self.ids[i][1] + n)):
                	# 如果是linux系统,下面这句代码中"\\"改为"/"
                    path_filename.append((self._imgpath % (self.ids[i][0], self.ids[i][1] + n)).split(self.root + "\\")[1])
                    break
# 第三处原始
        img = cv2.imread(self._imgpath % img_id, cv2.IMREAD_COLOR)
# 第三处修改
		img_names = ['.jpg','.png','.bmp', '.jpeg']
        for n in img_names:
            if os.path.exists(self._imgpath % img_id + n):
                tmp_imgPath = self._imgpath % img_id + n
                break
        img = cv2.imread(tmp_imgPath, cv2.IMREAD_COLOR)
  • 数据集格式

在这里插入图片描述

  • 训练命令
python tools/train.py -f exps/example/yolox_voc/yolox_voc_nano.py -d 0 -b 24 --fp16 -o -c yolox_nano.pth
  • 查看训练情况
cd YOLOX/YOLOX_outputs/yolox_voc_nano/tensorboard
tensorboard --logdir=./

在这里插入图片描述

二、转换模型

  1. 将pth转为onnx格式
python tools/export_onnx.py --output-name yolox_nano_self.onnx -f exps/example/yolox_voc/yolox_voc_nano.py -c YOLOX_outputs\yolox_voc_nano\best_ckpt.pth

则在项目目录下生成了yolox_nano_self.onnx的模型文件

  1. 参考博客基于NCNN将YOLOX部署到Android的2.3.2~2.3.4节将onnx转为bin和param格式

三、部署到rk3308

参考博客RK3308B部署mobilenetv2_ssdlite和使用YOLOX本地端调用USB摄像头实时检测将模型进行部署

更多推荐