YOLOX 训练自定义数据及海思平台适配

1.dataSet

1.folder

datasets|
  VOCdevkit|
    VOC2007|
      Annotations|
      ImageSets|
      JPEGImages|

2.split dataSet

import os
import random

trainval_percent = 0.1
train_percent = 0.9
xmlfilepath = 'Annotations'
txtsavepath = 'ImageSets/Main'
total_xml = os.listdir(xmlfilepath)


num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('ImageSets/Main/val.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        fval.write(name)
        ftest.write(name)
    else:
        ftrain.write(name)
ftrain.close()
fval.close()
ftest.close()

3.dataSet config

1. yolox/data/datasets/voc_classes.py，类别名
2. yolox/data/datasets/voc.py

    def __init__(
        self,
        data_dir,
        image_sets=[('2007', 'train')],
        img_size=(416,416), #修改为自己的size
        preproc=None,
        target_transform=AnnotationTransform(),
        dataset_name="VOC2007

2.model config

1.exp

模型中涉及的所有内容都放在一个单独的 Exp 文件中，包括模型设置、训练设置和测试设置。完整的 Exp 文件位于yolox_base.py。可能每个exp都写太长，但是你可以继承基础exp文件，只覆盖改变的部分。

我们还是以VOC Exp文件为例。

我们这里选择YOLOX-S模型，所以我们应该改变网络深度和宽度。VOC 只有 20 个类，因此我们也应该更改 num_classes。

这些配置在 init() 方法中更改：

class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.num_classes = 4
        self.depth = 0.33
        self.width = 0.375
        self.data_num_workers = 8
        self.input_size = (608, 608)
        self.scale = (0.5, 1.5)
        self.random_size = (10, 20)
        self.test_size = (608, 608)
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
        self.enable_mixup = False

2.NNIE

主要修改3个方面：

1. act loss
2. upsample
3. （optional）focus -> conv/reorg

此处主要针对act和updample进行修改：

backbone

backbone = YOLOPAFPN(self.depth, self.width, in_channels=in_channels, depthwise=True,act="lrelu")

• network_blocks.py

class DWConv(nn.Module):
    """Depthwise Conv + Conv"""
    def __init__(self, in_channels, out_channels, ksize, stride=1, act="lrelu"):

• yolo_pafpn.py

 # self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
self.upsample0 = nn.ConvTranspose2d(int(in_channels[1] * width),int(in_channels[1] * width),kernel_size=2, stride=2)
self.upsample1 = nn.ConvTranspose2d(int(in_channels[0] * width),int(in_channels[0] * width),kernel_size=2, stride=2)

head

head = YOLOXHead(self.num_classes, self.width, in_channels=in_channels, depthwise=True,act="lrelu")

• network_blocks.py

class DWConv(nn.Module):
    """Depthwise Conv + Conv"""
    def __init__(self, in_channels, out_channels, ksize, stride=1, act="lrelu"):

3.train

requirements

git clone git@github.com:Megvii-BaseDetection/YOLOX.git
cd YOLOX
pip3 install -U pip && pip3 install -r requirements.txt
pip3 install -v -e .  # or  python3 setup.py develop

#pycocotools
pip3 install cython; pip3 install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

#apex 加速torch训练
git clone https://github.com.cnpmjs.org/NVIDIA/apex.git
cd apex
pip3 install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

#apex错误，则：
cd apex
python setup.py install

单卡训练

python tools/train.py -f ./exps/electric/elect_tiny.py -d 1 -b 64 --fp16 -o -c /path/to/the/pretrained/weights

4.onnx

python tools/export_onnx.py --output-name yolox_s.onnx -n yolox-s -f exps/electric/elect_s.py -c YOLOX_outputs/elect_s/best_ckpt.pth.tar