训练FastestDet（Anchor-Free、参数量仅0.24M、计算量0.9GFLOPs），稍改代码使得符合YOLO数据集排布

0 参考链接

• 官方的代码：FastestDet

1 准备数据

我已有的数据集排布：（符合YOLO排布）

dataset:.
├─images
│  ├─test
│  │  ├─xxx.jpg
│  │  ├─xxx.jpg
│  ├─train
│  └─val
├─labels
│  ├─test
│  │  ├─xxx.txt
│  │  ├─xxx.txt
│  ├─train
│  └─val

1.1 使用以下代码生成绝对路径的txt文件

• root_dir：数据集根目录
• save_dir：存储目录
• 结果：生成3个txt文件abspath_xxx.txt

import os

if __name__ == '__main__':
    root_dir = r'F:/A_Publicdatasets/RDD2020-1202/train_valid/RDD2020_together/images'
    save_dir = 'F:/A_Publicdatasets/RDD2020-1202/train_valid/RDD2020_together/'

    for s in ['train', 'val', 'test']:
        save_path = f'{save_dir}/abspath_{s}.txt'
        with open(save_path, 'w') as f:
            for file in os.listdir(os.path.join(root_dir, s)):
                f.write(f'{root_dir}/{s}/{file}\n')

1.2 在config文件夹下新建一个xxx.names文件

例如，我新建了一个RDD2020.names

2 配置训练参数

拷贝configs/coco.yaml 文件并重取名，例如我重命名为RDD2020.yaml，然后进行以下改动：

3 稍改代码使得符合YOLO数据集排布

在utils/datasets.py中定位到def __getitem__(self, index)然后做如下改动：

label_path = img_path.replace('images', 'labels').replace('jpg', 'txt')

4 开始训练

我是直接在train.py里面设置了yaml位置，就可以直接运行train.py了

也可以不设置，然后调用如下参数在终端训练：

python train.py --yaml configs/coco.yaml

还可以进行以下改动，指明权重存放位置：

                save_dir = 'runs/'  # 存放训练文件的根目录
                save_name = 'FastestDet'    # 本次训练存放的文件名
                torch.save(self.model.state_dict(), f"{save_dir}/{save_name}/weight_AP05:%f_%d-epoch.pth" % (mAP05, epoch))
                # torch.save(self.model.state_dict(), f"checkpoint/weight_AP05:%f_%d-epoch.pth"%(mAP05, epoch))

5 使用eval.py进行验证

• 因为是进行验证，所以batchsize应该设置为1，则在configs/RDD2020.yaml中将BATCH_SIZE设置为1
• 如果想要用test来进行验证的话，就将configs/RDD2020.yaml中的VAL改为baspath_test.txt的路径

我喜欢直接在eval.py修改好了参数然后运行：

然后在终端输入python eval.py即可

如果想要下面打印的值的位数更多的话，可以：

1. 在eval.py中的最下面进行如下修改

    stats = evaluation.compute_map(val_dataloader, model)
    for element in stats:
        print(element)

2. 在utils/evaluation.py里面定位到coco_eval.summarize()然后进行如下修改：

        return coco_eval.stats

最终打印的结果就会多出最后面的那些值：

6 得到参数量、计算量的方法

TMD，终于试出来怎么打印计算量了！TNND！MLGBZ！

方式：

if __name__ == "__main__":
    model = FastestDet()
    # model.train() # 要打印下面参数量计算量的时候注释掉


    # -------------------打印参数量|计算量---------------- #
    model = model.model

    from thop import profile
    img = torch.rand(1, 3, 640, 640).to(device)
    flops, params = profile(model, (img,))
    print(f'params: {params} M, GFLOPs: {2 * flops} M')   # 注意这里的flops要×2，才跟yolo打印出来的值对应得上
    print('params: %.2f M, GFLOPs: %.1f M' % (params / 1e6, 2 * flops / 1e9))   # 注意这里的flops要×2，才跟yolo打印出来的值对应得上
    # -------------------打印参数量|计算量---------------- #