爆炸性新闻！女装大佬复现并深度优化YOLOX！

1 年前

学AI就上AI Studio，提供在线编程环境免费GPU算力

YOLOX

概述

女装大佬咩酱2021新作（隐瞒了大家这么久，其实我是一个女装大佬。希望能让大家认识一下另一个我！）！也是我以女装大佬身份的初投稿，大家点个喜欢和fork吧！我用飞桨深度学习框架复现了YOLOX算法并深度优化了其中的SimOTA部分！具体来讲，咩酱将SimOTA并行化！将YOLOX训练速度提高一个档次！原版YOLOX中使用for循环遍历每一张图片确定其正负样本，在dynamic_k_matching()中也使用for循环遍历每一个gt框以确定每一个gt框分配给几个预测框去学习。可见用了2层for循环。但是，这些在咩酱深度优化的SimOTA中统统都没有！咩酱干掉了所有for循环！一个批次的图片同时进行SimOTA操作！也不会遍历每一个gt！可以充分发挥飞桨深度学习框架的并行能力！我愿称之为飞桨python api的极致应用！代价也是有的，需要更多一些显存，即用空间换时间。

2021年，咩酱鸽了很久没有发精品项目，这不是缺席，而是希望项目在制作上更加精品，不停地打磨打磨，不鸣则已，一鸣惊人！整个项目耗费了我2个月的时间，是我最用心、完成度最高的项目。其中的视频剪辑耗费了我半个月的时间，看在我这么辛苦的份上，可不可以给咩酱的b站账号点个关注and给我的视频点个三连呢？视频文字版在“YOLOX训练、SimOTA优化详解(女装出镜)”那一章节。值得一提的是，咩酱在视频里设下两个彩蛋，分别是两大神作登场的名场面，至于是哪两大神作？卖个关子，你看到最后就知道了，一定要看到最后哦！初次见面，当然也有福利带给大家，本视频点赞过500，我为飞桨二创并演唱一首歌曲，本视频点赞每满1000，更有神秘的上不封顶的福利！赶快点赞和三连走起来吧！

我在 Paddle-YOLOX 仓库持续更新最新的代码，欢迎大家关注and下载代码到本机（windows系统也可以跑本项目）。为方便描述，下文中我将train_yolox.py脚本所在的目录称为YOLOX_HOME，比如，这个精品项目里，YOLOX_HOME为~/work；若各位读者用个人windows跑本项目，YOLOX_HOME为train_yolox.py脚本所在的目录。

一些结果：

注意: - 模型获取请见“快速开始”，权重移植到paddle后mAP有一点亏损，属正常现象； - 不测FPS（我要节能！），反正测了你也跑不过TensorRT，还不如把精力花在导出上； - test_dev费时费力（我要节能！），真想要测test_dev可以参考ppyolo的test_dev_yolo.py，只需改几处代码即可实现。

旷视提出了 YOLOX ，号称超越了一切的YOLO，是否是这样子呢？让咩酱来给大家深度剖析一下。让我们先来体验一下YOLOX的预测效果吧！

安装依赖

! cd ~/work; pip install -r requirements.txt

温馨提示：本项目涉及到的所有命令（包括训练、验证、预测）均在readme_yolox.txt里，大家可以打开这个文件并复制粘贴命令。

快速开始

(1)获取预训练模型

大家fork这个项目之后，到这里下载YOLOX的预训练模型，并上传到YOLOX_HOME目录下。这些预训练模型是怎么样得到的？和咩酱之前的精品项目一样，前往原版 YOLOX 仓库下载yolox_s.pth、yolox_m.pth、yolox_l.pth、yolox_x.pth这些预训练模型，再分别运行1_yolox_s_2paddle.py、1_yolox_m_2paddle.py、1_yolox_l_2paddle.py、1_yolox_x_2paddle.py转换得到。因为这些脚本里有import torch，读者可在自己的windows（或linux）上运行这些脚本，AIStudio上暂时不可。

(2)使用模型预测图片、获取FPS（预测images/test/里的图片，结果保存在images/res/）

! cd ~/work; python demo_yolox.py --config=0

--config=0表示使用了0号配置文件yolox_s.py，配置文件代号与配置文件的对应关系在tools/argparser.py文件里：

parser.add_argument('-c', '--config', type=int, default=0,
                            choices=[0, 1, 2, 3],
                            help=textwrap.dedent('''\
                            select one of these config files:
                            0 -- yolox_s.py
                            1 -- yolox_m.py
                            2 -- yolox_l.py
                            3 -- yolox_x.py'''))

即0代表使用config/yolox/yolox_s.py配置文件，1代表使用config/yolox/yolox_m.py配置文件，2代表使用config/yolox/yolox_l.py配置文件，3代表使用config/yolox/yolox_x.py配置文件。

预测时加载的权重为配置文件里test_cfg的model_path指定的文件，读者如果需要使用自己训练好的权重进行预测，需要修改model_path。

train_yolox.py、eval_yolox.py、demo_yolox.py都需要指定--config参数表示使用哪个配置文件，后面不再赘述。

数据集的放置位置

如果不是在AIStudio上训练，而是在个人电脑上训练，数据集应该和本项目位于同一级目录(同时需要修改一下配置文件中self.train_path、self.val_path这些参数使其指向数据集)。一个示例：

D://GitHub
     |------COCO
     |        |------annotations
     |        |------test2017
     |        |------train2017
     |        |------val2017
     |------VOCdevkit
     |        |------VOC2007
     |        |        |------Annotations
     |        |        |------ImageSets
     |        |        |------JPEGImages
     |        |        |------SegmentationClass
     |        |        |------SegmentationObject
     |        |
     |        |------VOC2012
     |                 |------Annotations
     |                 |------ImageSets
     |                 |------JPEGImages
     |                 |------SegmentationClass
     |                 |------SegmentationObject
     |------Paddle-PPYOLO-master
              |------annotation
              |------config
              |------data
              |------model
              |------...

训练

如果你需要训练COCO2017数据集，那么需要先解压数据集

! cd ~/data/data7122/; unzip ann*.zip
! cd ~/data/data7122/; unzip val*.zip
! cd ~/data/data7122/; unzip tes*.zip
! cd ~/data/data7122/; unzip image_info*.zip
! cd ~/data/data7122/; unzip train*.zip

根据自己的需要修改配置文件。以训练yolox_m为例，config/yolox/yolox_m.py配置文件部分内容如下：

self.train_cfg = dict(
            # batch_size=128,
            batch_size=8,
            num_workers=4,   # 读数据的进程数
            # model_path='dygraph_yolox_m.pdparams',
            model_path=None,
            # model_path='./weights/1000.pdparams',
            update_iter=1,    # 每隔几步更新一次参数
            log_iter=20,      # 每隔几步打印一次
            save_iter=1000,   # 每隔几步保存一次模型
            eval_epoch=10,    # 每隔几轮计算一次eval集的mAP。
            eval_iter=-1,    # 每隔几步计算一次eval集的mAP。设置了eval_epoch的话会自动计算。
            max_epoch=300,    # 训练多少轮
            max_iters=-1,     # 训练多少步。设置了max_epoch的话会自动计算。
            mosaic_epoch=285,  # 前几轮进行mosaic
            fp16=False,         # 是否用混合精度训练
            fleet=False,        # 是否用分布式训练
            find_unused_parameters=False,   # 是否在模型forward函数的返回值的所有张量中，遍历整个向后图。
        self.learningRate = dict(
            base_lr=0.01 * self.train_cfg['batch_size'] / 64,
            CosineDecay=dict(),
            LinearWarmup=dict(
                start_factor=0.,
                epochs=5,
        self.optimizerBuilder = dict(
            optimizer=dict(
                momentum=0.9,
                use_nesterov=True,
                type='Momentum',
            regularizer=dict(
                factor=0.0005,
                type='L2',
        )

batch_size我取8，这也是32GB V100能跑的不会爆显存的批大小；

num_workers为DataLoader需要的num_workers；

model_path=None表示不加载预训练权重，从头训练，和原版YOLOX一样；

update_iter=1表示每隔1步更新一次参数。或许读者可以尝试设置update_iter=8（基础学习率也要乘以8），即每隔8步更新一次参数，这样可以变相地增加了批大小=8*8=64，这样你就和8卡训练差不多了。值得一提的是，原版YOLOX并没有使用同步bn，而是用的普通bn，所以，大胆去尝试吧；

eval_epoch=10表示每隔10个epoch计算一次eval集的mAP；

max_epoch=300表示训练300个epoch，和原版YOLOX一样；

mosaic_epoch=285表示前285轮进行mosaic增强，和原版YOLOX一样；

fp16=False表示不使用混合精度训练，咩酱用COCO预训练模型迁移学习voc2012数据集时，发现用混合精度训练的话收敛会比较慢，而且有的梯度被裁剪为0，所以这里暂时先不用混合精度训练。或许也是咩酱对混合精度训练不熟悉，感兴趣的读者可以尝试调整GradScaler()的参数进行训练and在评论区留言提醒咩酱；

fleet=False表示不使用分布式训练。有多卡条件的读者可设置为True，并使用readme_yolox.txt里“分布式训练”处的命令启动训练（按需要修改--config=?）；

learningRate处使用了和原版YOLOX一样的配置，基础学习率base_lr=0.01 * self.train_cfg['batch_size'] / 64，前5个epoch学习率从0增加到基础学习率，后面学习率余弦衰减；

优化器的配置optimizerBuilder使用了和原版YOLOX一样的配置，使用Momentum优化器，滑动平均参数momentum=0.9，use_nesterov=True；网络中仅卷积层的weight参数使用L2衰减，factor=0.0005，bn层的weight、bias、卷积层的bias参数不使用L2衰减。

再输入以下命令训练（以训练yolox_m为例）

! cd ~/work; python train_yolox.py --config=1

训练5个epoch之后就能很好地检测到人。

注意，训练过程中有可能出现这个错误：

...
Error: /paddle/paddle/fluid/operators/bce_loss_op.cu:38 Assertion `(x >= static_cast<T>(0)) && (x <= one)` failed. Input is expected to be within the interval [0, 1], but recieved nan.
Error: /paddle/paddle/fluid/operators/bce_loss_op.cu:38 Assertion `(x >= static_cast<T>(0)) && (x <= one)` failed. Input is expected to be within the interval [0, 1], but recieved nan.
Error: /paddle/paddle/fluid/operators/bce_loss_op.cu:38 Assertion `(x >= static_cast<T>(0)) && (x <= one)` failed. Input is expected to be within the interval [0, 1], but recieved nan.
Error: /paddle/paddle/fluid/operators/bce_loss_op.cu:38 Assertion `(x >= static_cast<T>(0)) && (x <= one)` failed. Input is expected to be within the interval [0, 1], but recieved nan.
...

这个错误也困扰了咩酱很久。咩酱在训练yolox_s时会经常出现这个错误；训练yolox_m时则较少出现这个错误。所以咩酱才会用yolox_m作为示例。设置log_iter=1，即每一步都打印loss时，发现loss并不会出现nan。起初我以为是损失的问题，所以我将model/losses/iou_losses.py IOUloss()类的 call ()方法重写一遍，不使用逐句翻译的 call ()，后来我发现这样做只是会减少出现这个错误的概率（如果使用逐句翻译的 call ()，训练yolox_s(COCO2017数据集)几百几千步就出现这个错误，如果使用咩酱重写的 call ()，需要好几万步才出现这个错误）。咩酱走上了逐层排查的溯因之路，我每一步都打印loss和每一层的梯度，发现最开始是YOLOXHead的reg_preds[1].weight.grad、reg_preds[1].bias.grad中部分元素(梯度)为nan或者inf（当然，多跑几次发现reg_preds[0]、reg_preds[2]的梯度都有可能出现nan或者inf），导致再反向传播更新参数时将前面相关的层的权重污染为nan。导致下一次训练时，网络的输出为nan而出现这个错误。还有，我还发现，有可能reg_preds[0]、reg_preds[1]、reg_preds[2]的梯度都没有出现nan或者inf，而是稍微往前面一点的层，比如head.reg_convs[1][0].conv.weight.grad。总之， 错的不是我，不是分类分支，而是回归分支和整个世界！ 这个错误有可能在warmup阶段出现，所以调小学习率并不是解决这个错误的办法。

我与此错误周旋了很久，依然没有很好的解决方案。使用混合精度训练时，并不会出现这个错误，因为它天生就有防止浮点数溢出的特性。感兴趣的读者可以使用混合精度(可能需要调整一下GradScaler()的参数)训练从头训练coco并在评论区通知咩酱结果，以咩酱单线程实验验证，需要的时间太久。有好的解决方案也欢迎在评论区提出。如果不使用混合精度训练，咩酱给出的解决方案是，鸵鸟算法！也就是出现错误并中断训练之后，手动修改配置文件中train_cfg的model_path='./weights/xxx.pdparams'，xxx表示保存的最新的模型。再运行以上命令接着训练。一个很好的建议是编写一个脚本(python或shell)，不停地监控训练进程，当监测到训练进程死亡之后，自动修改配置文件中train_cfg的model_path为最新，然后重启训练命令。我为什么不写？因为我的耐心被耗完了。

一个比较残酷的事实是，当出现上面这个错误之后，接着训练就会很容易再出现这个错误，导致训练无法再进行下去，而且验证集mAP直接变成0。所以是很难训练COCO数据集的！所以我推荐大家用这个仓库来迁移学习一下自定义数据集就好。另外我也希望能有大佬帮助我解决这个错误！现在咩酱正在尝试使用混合精度训练，有结果会通知大家。

训练自定义数据集（迁移学习）

自带的voc2012数据集是一个很好的例子。

先解压voc数据集：

! cd ~/data/data4379/; unzip pascalvoc.zip

将自己数据集的txt注解文件放到annotation目录下，txt注解文件的格式如下：

xxx.jpg 18.19,6.32,424.13,421.83,20 323.86,2.65,640.0,421.94,20
xxx.jpg 48,240,195,371,11 8,12,352,498,14
# 图片名 物体1左上角x坐标,物体1左上角y坐标,物体1右下角x坐标,物体1右下角y坐标,物体1类别id 物体2左上角x坐标,物体2左上角y坐标,物体2右下角x坐标,物体2右下角y坐标,物体2类别id ...

注意：xxx.jpg仅仅是文件名而不是文件的路径！xxx.jpg仅仅是文件名而不是文件的路径！xxx.jpg仅仅是文件名而不是文件的路径！

运行1_txt2json.py会在annotation_json目录下生成两个coco注解风格的json注解文件(如果是其它数据集，你还需要修改一下1_txt2json.py中的train_path、val_path、test_path、classes_path、train_pre_path、val_pre_path、test_pre_path)，这是训练脚本train_yolox.py支持的注解文件格式。

还是以yolox_m为例，在config/yolox/yolox_m.py里修改train_path、val_path、classes_path、train_pre_path、val_pre_path、num_classes这6个变量（自带的voc2012数据集直接解除注释就ok了）,就可以开始训练自己的数据集了。

另外，根据自己的需求修改配置文件里的其它配置

self.train_cfg = dict(
            # batch_size=128,
            batch_size=8,
            num_workers=4,   # 读数据的进程数
            model_path='dygraph_yolox_m.pdparams',
            # model_path=None,
            # model_path='./weights/1000.pdparams',
            update_iter=1,    # 每隔几步更新一次参数
            log_iter=20,      # 每隔几步打印一次
            save_iter=500,   # 每隔几步保存一次模型
            eval_epoch=1,    # 每隔几轮计算一次eval集的mAP。
            eval_iter=-1,    # 每隔几步计算一次eval集的mAP。设置了eval_epoch的话会自动计算。
            max_epoch=4,    # 训练多少轮
            max_iters=-1,     # 训练多少步。设置了max_epoch的话会自动计算。
            mosaic_epoch=3,  # 前几轮进行mosaic
            fp16=False,         # 是否用混合精度训练
            fleet=False,        # 是否用分布式训练
            find_unused_parameters=False,   # 是否在模型forward函数的返回值的所有张量中，遍历整个向后图。
        self.learningRate = dict(
            base_lr=0.01 * self.train_cfg['batch_size'] / 64,
            CosineDecay=dict(),
            LinearWarmup=dict(
                start_factor=0.,
                epochs=1,
        self.backbone = dict(
            dep_mul=0.67,
            wid_mul=0.75,
            freeze_at=5,
            fix_bn_mean_var_at=0,
        )

model_path='dygraph_yolox_m.pdparams'表示读取预训练模型'dygraph_yolox_m.pdparams'进行训练；

freeze_at=5表示冻结骨干网络进行训练，这样可以减少显存需求以及加快训练速度。

还是使用同样的命令启动训练：

! cd ~/work; python train_yolox.py --config=1

实测yolox_m的AP(0.50:0.95)可以到达0.62+、AP(small)可以到达0.25+。

demo_yolox.py、eval_yolox.py也是根据配置文件指定的数据集进行预测、验证。

评估

运行以下命令。评测的模型是config/yolox/yolox_m.py里self.eval_cfg -> model_path指定的模型

! cd ~/work; python eval_yolox.py --config=1

该mAP是val集的结果。

预测

运行以下命令。使用的模型是config/yolox/yolox_m.py里self.test_cfg -> model_path指定的模型

! cd ~/work; python demo_yolox.py --config=1

YOLOX训练、SimOTA优化详解(女装出镜)

终于到了压轴环节，这是咩酱之前发的精品项目里没有的环节。以前我总是只想写代码，写完只剩半口气了，哪还有精力分享技术细节。但是，精品项目不能只有这些，一个好的精品项目，应该对每一位读者负全责到底！我用飞桨深度学习框架复现了YOLOX算法并深度优化了其中的SimOTA部分！具体来讲，咩酱将SimOTA并行化！将YOLOX训练速度提高一个档次！原版YOLOX中使用for循环遍历每一张图片确定其正负样本，在dynamic_k_matching()中也使用for循环遍历每一个gt框以确定每一个gt框分配给几个预测框去学习。可见用了2层for循环。但是，这些在咩酱深度优化的SimOTA中统统都没有！咩酱干掉了所有for循环！一个批次的图片同时进行SimOTA操作！也不会遍历每一个gt！可以充分发挥飞桨深度学习框架的并行能力！我愿称之为飞桨python api的极致应用！代价也是有的，需要更多一些显存，即用空间换时间。

讲了半天，我好累，下面有请另一个我来给大家介绍这一部分吧！ 拜托了！另一个我！

是是是！ 哈喽！大家好，初次见面！我是寄宿在咩酱体内的另一个人格，大家可以叫我糖妹，以和咩酱区分。糖妹特地做了一个视频详解这一部分，即“概述”处的视频。希望大家能给糖妹点个关注and给个三连！你们的三连对我非常重要！

同时，糖妹会在AIStudio和大家同步播报文字版。糖妹来给大家详解YOLOX的训练部分。

YOLOv3前置知识

首先，我们先来说说YOLOv3算法吧。众所周知，一张图片经过YOLOv3网络前向传播之后，会输出3个不同分辨率的特征图。比如输入网络的图片形状为(1, 3, 416, 416)时，经过网络前向传播后，会得到形状分别为(1, 255, 13, 13)、(1, 255, 26, 26)、(1, 255, 52, 52)的3个特征图（数据集为COCO时）。如图1所示：

(图1)

我们把特征图和原图对齐来看的话，原图好像被切割成了一个一个的格子！一个格子对应特征图的一个像素点。我们可以大致看到每个格子大概拥有原图哪些区域的特征。需要注意的是，因为图片经过多层卷积层，叠加的感受野是非常大的，所以每个格子不仅仅只有格子内部原图区域的信息，还包括格子之外的感受野可达的原图区域的信息。

每个特征图有255个通道，这是因为，每个特征图的一个像素点（即一个格子）会出3个预测框，每个预测框有85位信息，0到3位是预测框的中心点xy偏移和预测框的宽高，第4位是objness，后80位是80个类别的条件概率。所以每个像素应该有3 * 85=255位信息。我们发现，特征图的分辨率比输入图片的分辨率小，这是因为网络中存在步长为2的卷积层，图片经过它，会使得自己的宽高缩小一半（即下采样）。(1, 255, 13, 13)形状的特征图经过了5次步长为2的卷积层，所以宽高变成了416/2^5=13；(1, 255, 26, 26)形状的特征图经过了4次步长为2的卷积层，所以宽高变成了416/2^4=26；(1, 255, 52, 52)形状的特征图经过了3次步长为2的卷积层，所以宽高变成了416/2^3=52。当然，这些特征图也不是输入图片一路卷到底来的，YOLOv3网络中有一个类似FPN的结构，低分辨率的特征图会进行一个上采样操作（最近邻插值）和高分辨率的特征图进行通道维的concat来提高检测效果，如图2所示：

(图2)

图中Output表示的是特征图的大小，图中的输入图片是256 * 256的大小。右下角斜着的那两条线就代表上采样了，用的是最近邻插值。汇合时的小黑点代表concat。

另外，下文中我还会经常提到一个词，特征图的stride，即特征图的步长。比如形状为(1, 255, 13, 13)的特征图，它经过了5次步长为2的卷积层，所以它的stride=2^5=32。可以理解为特征图的一个像素其实跨过了输入图片的32个像素。更形象地，你也可以理解为“格子的边长”，如图1右上角的特征图所示，13x13的特征图和416x416的输入图片对齐之后，就好像输入图片被切割成13x13个格子，每个格子的边长为32。同理，图1右下角的特征图的stride=16，图1左下角的特征图的stride=8。

YOLOX与YOLOv3的一些不同

好的，说了这么多铺垫终于轮到YOLOX出场了！YOLOX预测时默认使用了640x640的分辨率，直观一点，我再画一张图：

(图3)

YOLOX与YOLOv3有相似之处，也有不同之处。它也是像YOLOv3一样，一张图片出3种感受野的预测。不同的是，它使用了解耦头，即对于每种感受野，使用不同的卷积层作用于特征图以分别预测预测框的xy宽高、objness、clsness。所以每种感受野出3个特征图，通道数分别为4、1、80。

还有，YOLOX是一个Anchor-Free算法，而且每个格子只会出1个预测框，而YOLOv3是每个格子出3个预测框。还有，它的网络结构比YOLOv3的网络结构要复杂（本视频不会涉及过多的网络结构细节）。

废话不多说，我们开始训练YOLOX吧！

YOLOX的数据预处理

训练第一步当然是从预处理一小批图片开始。train_yolox.py中COCOTrainDataset类中的 getitem (self, idx)方法用来读取一张图片。图片会经过预处理，在配置文件config/yolox/yolox_s.py中有预处理相关的配置，我们看代码，预处理部分的代码是借鉴了PaddleDetection的，所以，熟悉PaddleDetection的小伙伴应该很容易看懂。

# config/yolox/yolox_s.py
        # DecodeImage
        self.decodeImage = dict(
        # YOLOXMosaicImage
        self.yOLOXMosaicImage = dict(
        # ColorDistort
        self.colorDistort = dict()
        # RandomFlipImage
        self.randomFlipImage = dict(
            is_normalized=False,
        # BboxXYXY2XYWH
        self.bboxXYXY2XYWH = dict()
        # YOLOXResizeImage
        self.yOLOXResizeImage = dict(
            target_size=[640 - i * 32 for i in range(7)] + [640 + i * 32 for i in range(1, 7)],
            # target_size=[640],
            interp=1,  # cv2.INTER_LINEAR = 1
            use_cv2=True,
            resize_box=True,
        # PadBox
        self.padBox = dict(
            num_max_boxes=120,
            init_bbox=[-9999.0, -9999.0, 10.0, 10.0],
        # SquareImage
        self.squareImage = dict(
            fill_value=114,
            is_channel_first=False,
        # Permute
        self.permute = dict(
            to_bgr=False,
            channel_first=True,
        # 预处理顺序。增加一些数据增强时这里也要加上，否则train.py中相当于没加！
        self.sample_transforms_seq = []
        self.sample_transforms_seq.append('decodeImage')
        self.sample_transforms_seq.append('yOLOXMosaicImage')
        self.sample_transforms_seq.append('colorDistort')
        self.sample_transforms_seq.append('randomFlipImage')
        self.sample_transforms_seq.append('bboxXYXY2XYWH')
        self.sample_transforms_seq.append('yOLOXResizeImage')
        self.sample_transforms_seq.append('padBox')
        self.batch_transforms_seq = []
        self.batch_transforms_seq.append('squareImage')
        self.batch_transforms_seq.append('permute')
...

首先经过DecodeImage打开（解码）图片，然后经过马赛克增强、色彩扭曲、随机水平翻转这些数据增强，之后经过BboxXYXY2XYWH，也就是将gt框的格式从“左上角xy坐标+右下角xy坐标”变成“gt框中心点xy坐标+gt框宽高”。YOLOXResizeImage的作用是随机选一个尺度进行Resize，即多尺度训练，这一步gt框也会跟着图片一起缩放，而且图片缩放时是保持着图片原始宽高比进行缩放。PadBox的作用是将每张图片的gt数量填充到num_max_boxes=120个，gt_class和gt_score也会跟着填充到120个。这么做的原因是为了可以用一个形状为(batch_size, 120, 4)的张量表示这一批图片的所有gt框的xywh信息；init_bbox=[-9999.0, -9999.0, 10.0, 10.0]表示填充的gt框的信息初始为这4个数，为什么用这些数字初始化？一会讲到SimOTA时会与大家说。SquareImage的作用是将图片填充成正方形，fill_value=114表示填充的颜色是灰色。刚才说到，YOLOXResizeImage这一步缩放时是保持着图片原始宽高比进行缩放的，假如我一张宽240高320的图片resize到640这个尺度，那么宽会变成480，高会变成640。那么我就需要SquareImage这一步将图片填充成一个正方形图片。SquareImage的代码也非常简单，is_channel_first=True表示输入的图片维度是CHW排序的，is_channel_first=False表示输入的图片维度是HWC排序的；先创建一张边长为max(H, W)的图片padded_img，这里为640，然后将im左上角对准padded_img左上角贴上去就完事了。这一步不会改动gt的坐标，所以gt不用跟着变换一下。

# tools/transform.py
class SquareImage(BaseOperator):
    def __call__(self, sample, context=None):
                    im = sample[k]
                    im = im.astype(np.float32, copy=False)
                    sample[k] = im  # 类型转换
                    if self.is_channel_first:
                        C, H, W = im.shape
                        if H != W:
                            max_ = max(H, W)
                            padded_img = np.ones((C, max_, max_), dtype=np.uint8) * self.fill_value
                            padded_img = padded_img.astype(np.float32)
                            padded_img[:C, :H, :W] = im
                            sample[k] = padded_img
                    else:
                        H, W, C = im.shape
                        if H != W:
                            max_ = max(H, W)
                            padded_img = np.ones((max_, max_, C), dtype=np.uint8) * self.fill_value
                            padded_img = padded_img.astype(np.float32)
                            padded_img[:H, :W, :C] = im
                            sample[k] = padded_img
                    break
        if not batch_input:
            samples = samples[0]
        return samples
...

Permute的作用是将图片维度顺序从HWC变成CHW。

预处理完这一批图片之后，我们回到train_yolox.py

# train_yolox.py
class COCOTrainDataset(paddle.io.Dataset):
    def __init__(self, records, init_iter_id, cfg, sample_transforms, batch_transforms):
    def __getitem__(self, idx):
        # 取出感兴趣的项
        image = sample['image'].astype(np.float32)
        gt_bbox = sample['gt_bbox'].astype(np.float32)
        # gt_score = sample['gt_score'].astype(np.float32)
        gt_class = sample['gt_class'].astype(np.int32)
        gt_class = np.expand_dims(gt_class, 1).astype(np.float32)
        gt_class_bbox = np.concatenate([gt_class, gt_bbox], 1)
        return image, gt_class_bbox
    def __len__(self):
        return len(self.indexes)