训练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即可

如果想要下面打印的值的位数更多的话,可以:
- 在
eval.py中的最下面进行如下修改

stats = evaluation.compute_map(val_dataloader, model)
for element in stats:
print(element)
- 在
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打印出来的值对应得上
# -------------------打印参数量|计算量---------------- #


DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐

所有评论(0)