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

发布时间:2024年01月14日

0 参考链接

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))
文章来源:https://blog.csdn.net/LWD19981223/article/details/135589153
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。