项目场景:
在YOLOv9
模型的改进中,常常需要替换一些模块来提高模型的精度。但在评价模型大小规模的时候需要根据模型的参数量、计算量进行评定,一般在模型的训练文件train.py
,或者是test.py
还有models/yolo.py
都会输出这些数据。
问题描述
在模型改进过程中,在YOLOv9
的基础模型中添加了新模块,在训练时,打印店信息中没有打印参数量GFLOPs
。在YOLOv7
的改进中,未打印GFLOPs
的只需要安装thop
包就可以正常打印了。但在YOLOv9
中似乎不行。例如此处运行models/yolo.py
,未打印GFLOPs
。
解决方案:
在models/yolo.py
的if __name__ == '__main__'
下方添加如下代码,自行输出:
from thop import profile
img = torch.rand(1, 3, 640, 640).to(device)
flops, params = profile(model, (img,))
print('parameters: %.2f M, GFLOPs: %.1f GFLOPs' % (params / 1e6, 2 * flops / 1e9))
具体添加位置:
运行结果:
此时就已经输出GFLOPs
了,并且输出的参数量也是一致的。