使用YOLOv6项目训练数据,默认没有计算FLOPS的选项,这时可以通过引入fvcore来计算FLOPS
Github地址
- 安装fvcore
进入python环境,安装fvcore
pip install fvcore
-
找到eval.py文件
位于tools/eval.py -
导入fvcore
在eval.py文件中导入
from fvcore.nn import FlopCountAnalysis, parameter_count_table
-
找到插入代码的位置
-
在上图代码的下方插入以下代码
pbar = enumerate(dataloader)
for i, (imgs, targets, paths, _) in pbar: # batch
imgs = imgs.to(device, non_blocking=True).float()
f1 = FlopCountAnalysis(model, imgs)
print('FLPOS:', f1.total())
- 修改eval-interval方便查看效果(非必要)
接下来可以去train中修改eval-interval参数改为1,这样就可以让模型在训练一轮后就输出计算的参数
parser.add_argument('--eval-interval', default=1, type=int, help='evaluate at every interval epochs')
- 完工
接下来启动训练train.py,就可以看到输出的FLOPS参数啦