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

news2024/11/20 15:29:40

文章目录

0 参考链接
1 准备数据
- 1.1 使用以下代码生成绝对路径的txt文件
- 1.2 在config文件夹下新建一个xxx.names文件
2 配置训练参数
3 稍改代码使得符合YOLO数据集排布
4 开始训练

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))