项目场景:
在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了,并且输出的参数量也是一致的。


















