对yolox框架进行自定义改进时遇到的问题及解决记录

-------先写一部分，后续更新--------

-------2022.07.16 更新yolox训练过程中验证集评估时出现KeyError错误的分析及解决办法

-------2022.07.28 更新yolox如何自定义从backbone传入neck的特征金字塔部分的三个层次的channel通道数大小

前言

需要注意的：

yolox代码（这里指其团队在github上开源的代码）体量大结构复杂，虽然能看出来开发者的代码编写和模块安排还是很规范、合理的，但由于各功能模块太多，在进行训练时仍要层层迭代调用各文件夹、python文件、类、函数，这就给程序理解以及后期自己在yolox上进行改进带来的一定的麻烦。

（1）依照官方提供的README文件的说明，我们从终端窗口运行tools/train.py文件开始训练，一切的训练调用都是从这个文件为根基发散的。

（2）Exp类的定义调用是yolox训练的核心，yolox将训练的具体配置、以及在训练过程中的某些函数统一打包在了这个Exp类里面。在运行train.py时，要将准备使用的Exp类所在的python文件地址作为参数传入。

例如，我在使用voc格式的数据集时，就是直接基于yolox源码提供的yolox_voc_s.py文件的基础上修改来作为我用来训练的Exp类。尤其要注意的是，yolox_voc_s.py文件中定义的Exp类并不是直接一步到位完成定义的，它是在定义时继承了yolox/exp/yolox_base.py中定义的Exp类（yolox_base.py中定义的Exp类又是继承的yolox/exp/base_exp.py中的BaseExp类，但是Exp类的大部分属性、类函数都是在yolox_base.py中定义的Exp类中实现的）。因此虽然我们表面上只用了yolox_voc_s.py文件中定义的Exp类，但是如果要仔细修改Exp类的内容我们常常是需要追溯到它继承的yolox/exp/yolox_base.py中定义的Exp类来进行修改的。要特别注意的是，在程序运行、修改过程中也不要盯着yolox_base.py中的Exp类，因为在yolox_voc_s.py中的Exp类在继承时可能会对yolox_base.py中的Exp类的个别类属性、类函数进行覆盖修改，只看yolox_base.py中的Exp类的类属性参数作参考可能会发生误判。（例如我因为只看着yolox_base.py中的Exp类中定义的self.depth、self.width这两个类属性怎么也排查不了bug所在，结果后来发现这两个类属性的默认值在yolox_voc_s.py中的Exp类继承后修改了，这才导致我之前反复核对也没发现，所以这一点要特别注意。）

1.模型定义文件介绍：

yolox的模块化分的比较清，yolox模型大体上就开看作backbone、以fpn为主的neck、head。这里yolox在搭建模型时，模型的主模块都放在了yolox/models文件夹下：

darknet.py定义了backbone；

yolo_pafpn定义了带pan模块的fpn结构作为模型的neck，同时在其中定义的YOLOPAFPN类中的其实还用self.backbone这个类属性实例化了darknet.py定义的模型backbone，并在forward函数中运算了backbone，并对backbone输出结果与其设计的fpn连接。因此yolo_pafpn实际上是包括了yolox框架的backbone和fpn，后面只要再加上head部分即可构成完整网络。

yolo_head.py定义了YOLOXHead类，里面打包了YOLOX的head部分，在代码体量上也是模型中最长、最复杂的部分。

yolox.py主要是将上面介绍的YOLOPAFPN类和YOLOXHead类包含的结构组成yolox模型，并对模型输出结果整理出包含各部分数值的字典。

2.模型改进注意点：

backbone后面分为dark2、dark3、dark4、dark5这几个线性运算模块，其中dark3、dark4、dark5后面的计算结果要分别保存，传到后面fpn的三个层次中进行计算。因此，在对模型backbone进行改动时，要特别注意dark3、dark4、dark5计算后的图像通道数能否与其要出入的fpn对应层次运算模块的实现准备接受的图像通道数一致。（这里再次强调，在exp类中的self.depth、self.width两个类属性会影响模型深度（总层数）和模型对图像运算时设定的通道数）

3.yolox训练过程中验证集评估时出现KeyError错误的分析及解决办法

在yolox训练中，在规定的一些epoch中，模型在训练集训练后会在验证集上进行验证测试性能，我在对训练集进行更换后出现了如下报错提示：

07-16 13:56:51 | INFO     | yolox.core.trainer:266 - epoch: 1/300, iter: 110/116, mem: 5523Mb, iter_time: 0.229s, data_time: 0.000s, total_loss: 10.5, iou_loss: 3.8, l1_loss: 0.0, conf_loss: 5.7, cls_loss: 0.9, lr: 1.124e-03, size: 640, ETA: 3:04:37
2022-07-16 13:56:52 | INFO     | yolox.core.trainer:370 - Save weights to ./YOLOX_outputs/yolox_voc_s
100%|##########| 29/29 [00:02<00:00, 10.64it/s]
2022-07-16 13:56:55 | INFO     | yolox.evaluators.voc_evaluator:160 - Evaluate in main process...
Writing ship VOC results file
Eval IoU : 0.50
2022-07-16 13:56:58 | INFO     | yolox.core.trainer:208 - Training of experiment is done and the best AP is 0.00， the best AP_50 is 0.00

2022-07-16 13:56:58 | ERROR    | yolox.core.launch:98 - An error has been caught in function 'launch', process 'MainProcess' (25844), thread 'MainThread' (140597191565888):
Traceback (most recent call last):

  File "tools/train.py", line 151, in <module>
    launch(
    └ <function launch at 0x7fde2b82af70>

> File "/home/dwt/MyCode/pycharm_projects/YOLOX_sample/yolox/core/launch.py", line 98, in launch
    main_func(*args)
    │          └ (╒═══════════════════╤═══════════════════════════════════════════════════════════════════════════════════════════════════════...
    └ <function main at 0x7fde1d3d4280>

  File "tools/train.py", line 134, in main
    trainer.train()#这里运行yolox/core下的trainer.py文件中的类Trainer的一个实例trainer下的类Trainer中定义的类函数train()，其中类Trainer通过定义各种类函数来实现了模型训练过程
    │       └ <function Trainer.train at 0x7fde18897940>
    └ <yolox.core.trainer.Trainer object at 0x7fde188a60a0>

  File "/home/dwt/MyCode/pycharm_projects/YOLOX_sample/yolox/core/trainer.py", line 88, in train
    self.train_in_epoch()
    │    └ <function Trainer.train_in_epoch at 0x7fde1889f040>
    └ <yolox.core.trainer.Trainer object at 0x7fde188a60a0>

  File "/home/dwt/MyCode/pycharm_projects/YOLOX_sample/yolox/core/trainer.py", line 98, in train_in_epoch
    self.after_epoch()
    │    └ <function Trainer.after_epoch at 0x7fde1889f3a0>
    └ <yolox.core.trainer.Trainer object at 0x7fde188a60a0>

  File "/home/dwt/MyCode/pycharm_projects/YOLOX_sample/yolox/core/trainer.py", line 235, in after_epoch
    self.evaluate_and_save_model()
    │    └ <function Trainer.evaluate_and_save_model at 0x7fde1889f670>
    └ <yolox.core.trainer.Trainer object at 0x7fde188a60a0>

  File "/home/dwt/MyCode/pycharm_projects/YOLOX_sample/yolox/core/trainer.py", line 339, in evaluate_and_save_model
    ap50_95, ap50, summary = self.exp.eval(
                             │    │   └ <function Exp.eval at 0x7fde18897f70>
                             │    └ ╒═══════════════════╤════════════════════════════════════════════════════════════════════════════════════════════════════════...
                             └ <yolox.core.trainer.Trainer object at 0x7fde188a60a0>

  File "/home/dwt/MyCode/pycharm_projects/YOLOX_sample/yolox/exp/yolox_base.py", line 335, in eval
    return evaluator.evaluate(model, is_distributed, half)
           │         │        │      │               └ False
           │         │        │      └ False
           │         │        └ fbnet_YOLOX(
           │         │            (backbone): fbnet_YOLOPAFPN(
           │         │              (backbone): fbnet_backbone(
           │         │                (stem): Focus(
           │         │                  (conv): BaseConv(
           │         │            ...
           │         └ <function VOCEvaluator.evaluate at 0x7fde188914c0>
           └ <yolox.evaluators.voc_evaluator.VOCEvaluator object at 0x7fddfc162a00>

  File "/home/dwt/MyCode/pycharm_projects/YOLOX_sample/yolox/evaluators/voc_evaluator.py", line 128, in evaluate
    eval_results = self.evaluate_prediction(data_list, statistics)
                   │    │                   │          └ tensor([ 2.2012,  0.0739, 28.0000], device='cuda:0')
                   │    │                   └ {0: (tensor([[-7.5409e+01, -6.3225e+01,  5.1623e+02,  2.6991e+02],
                   │    │                             [-1.3955e+02, -6.8198e+00,  5.6981e+02,  2.1258e+0...
                   │    └ <function VOCEvaluator.evaluate_prediction at 0x7fde188915e0>
                   └ <yolox.evaluators.voc_evaluator.VOCEvaluator object at 0x7fddfc162a00>

  File "/home/dwt/MyCode/pycharm_projects/YOLOX_sample/yolox/evaluators/voc_evaluator.py", line 205, in evaluate_prediction
    mAP50, mAP70 = self.dataloader.dataset.evaluate_detections(
                   │    │          │       └ <function VOCDetection.evaluate_detections at 0x7fde18891e50>
                   │    │          └ <yolox.data.datasets.voc.VOCDetection object at 0x7fde188b6df0>
                   │    └ <torch.utils.data.dataloader.DataLoader object at 0x7fddfc162dc0>
                   └ <yolox.evaluators.voc_evaluator.VOCEvaluator object at 0x7fddfc162a00>

  File "/home/dwt/MyCode/pycharm_projects/YOLOX_sample/yolox/data/datasets/voc.py", line 271, in evaluate_detections
    mAP = self._do_python_eval(output_dir, iou)
          │    │               │           └ 0.5
          │    │               └ '/tmp/tmpiht1ubmu'
          │    └ <function VOCDetection._do_python_eval at 0x7fde18897040>
          └ <yolox.data.datasets.voc.VOCDetection object at 0x7fde188b6df0>

  File "/home/dwt/MyCode/pycharm_projects/YOLOX_sample/yolox/data/datasets/voc.py", line 337, in _do_python_eval
    rec, prec, ap = voc_eval(
                    └ <function voc_eval at 0x7fde188913a0>

  File "/home/dwt/MyCode/pycharm_projects/YOLOX_sample/yolox/evaluators/voc_eval.py", line 109, in voc_eval
    R = [obj for obj in recs[imagename] if obj["name"] == classname]
                        │    │                            └ 'ship'
                        │    └ '000033'
                        └ {}

KeyError: '000033'
(fbnet_yolox) dwt@dwt-Ubuntu:~/MyCode/pycharm_projects/YOLOX_sample$ python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 1 -b 8 --fp16 -o
2022-07-16 13:59:59 | INFO     | yolox.core.trainer:142 - args: Namespace(batch_size=8, cache=False, ckpt=None, devices=1, dist_backend='nccl', dist_url=None, exp_file='exps/example/yolox_voc/yolox_voc_s.py', experiment_name='yolox_voc_s', fp16=True, logger='tensorboard', machine_rank=0, name=None, num_machines=1, occupy=True, opts=[], resume=False, s

初步判断应该是在读取验证集文件时出现错误，通过搜索，发现有类似问题的博客：yolov5 v3.0训练出现KeyError错误

参考该博客内容得知：

出现这个错误的原因是由于，数据的缓存文件可能发生了变化，yolov5默认会把读取的数据保存到缓存文件中，这样下载读取的速度就会很快！

类似的，我通过查看yolox中的voc_eval.py文件源代码：

def voc_eval(
    detpath,
    annopath,
    imagesetfile,
    classname,
    cachedir,
    ovthresh=0.5,
    use_07_metric=False,
):
    # first load gt
    if not os.path.isdir(cachedir):
        os.mkdir(cachedir)
    cachefile = os.path.join(cachedir, "annots.pkl")
    # read list of images
    with open(imagesetfile, "r") as f:
        lines = f.readlines()
    imagenames = [x.strip() for x in lines]

    if not os.path.isfile(cachefile):
        # load annots
        recs = {}
        for i, imagename in enumerate(imagenames):
            recs[imagename] = parse_rec(annopath.format(imagename))
            if i % 100 == 0:
                print("Reading annotation for {:d}/{:d}".format(i + 1, len(imagenames)))
        # save
        print("Saving cached annotations to {:s}".format(cachefile))
        with open(cachefile, "wb") as f:
            pickle.dump(recs, f)
    else:
        # load
        with open(cachefile, "rb") as f:
            recs = pickle.load(f)

从"cachedir"等字样猜测yolox应该也是有类似的对数据缓存的功能以加快后面的数据读取速度，因此本次报错可能是由于yolox还使用我之前用的数据集缓存与我更换后的数据集对应不上导致的。由此找到了缓存文件annots.pkl文件的所在地址：datasets/VOCdevkit/annotations_cache直接将该文件夹删除（yolox再次运行时如果找不到该文件夹会再次新建重新缓存数据），再次运行训练命令开始训练后该问题没有再次出现。

4.从backbone传入neck的特征金字塔部分的三个层次的channel通道数大小

阅读yolox源码，从定义yolox框架的neck模块的文件 yolox/models/yolo_pafpn.py中的源码部分（如下图所示）：

 可以看出其中的in_channel = [256, 512, 1024]参数，说明了yolox默认是让neck部分的特征金字塔依次接收256、512、1024这三种通道数的feature map特征图。当然，在该部分代码还可以看见有一个width参数，这个参数在后面部分与in_channel中的数字相乘，一定程度上调节了特征金字塔接收的特征图的通道数尺寸（如下图所示）：

 例如，原定的in_channel参数为  [256, 512, 1024]，而当width等于0.5时，则通过width * in_channel 可以将特征金字塔实际接收的特征图尺寸变为 [128,256,512]。但是仅仅通过调节width虽然可以对 in_channel起到一定改变，但特征金字塔三个层次的接收的特征图通道尺寸的比例是不变的，一直都是1:2:4。这样对通道数方面的调节还是不够灵活，所以我们要直接从in_channel上动手修改。

直接在yolox/models/yolo_pafpn.py中的类YOLOPAFPN中修改参数in_channels 和width是行不通的，因为这里只是yolox的neck部分的定义文件，这里定义的in_channels 和width只是默认数值，实际运行训练时，会在调用这一部分定义时给参数in_channels 和width传入新值。所以我们要找到调用neck部分来修改参数in_channels 和width。在yolox的默认训练文件下，通过调用yolox模型框架各模块文件来组装yolox模型进而训练的函数是定义在 yolox/exp/yolox_base.py中的类函数 get_model()部分。如下图所示：

可以看到这里调用了YOLOPAFPN，传入了新的in_channels参数，因此我们在这里修改的 in_channels参数值才会真正改变实际训练中yolox模型的特征金字塔的输入特征图通道数。而width参数的修改要看你是调用哪个文件中定义的EXP类，如果是调用exps/example/yolox_voc/yolox_voc_s.py中的EXP类，则要到该文件中找到width参数修改。具体原因可以参见本文章上面前言部分内容。