我已有的数据集排布:(符合YOLO排布)
dataset:.
├─images
│ ├─test
│ │ ├─xxx.jpg
│ │ ├─xxx.jpg
│ ├─train
│ └─val
├─labels
│ ├─test
│ │ ├─xxx.txt
│ │ ├─xxx.txt
│ ├─train
│ └─val
root_dir
:数据集根目录save_dir
:存储目录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')
例如,我新建了一个RDD2020.names
拷贝configs/coco.yaml
文件并重取名,例如我重命名为RDD2020.yaml
,然后进行以下改动:
在utils/datasets.py
中定位到def __getitem__(self, index)
然后做如下改动:
label_path = img_path.replace('images', 'labels').replace('jpg', 'txt')
我是直接在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))