文章目录
- 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 osif __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))