使用mmengine.analysis.get_model_complexity_info()打印任意模型参数和flops

原创

已于 2024-03-31 15:09:56 修改 · 1.8k 阅读

4 ·

CC 4.0 BY-SA版权

本文为博主原创文章，未经博主允许不得转载

文章标签：

#python #深度学习 #目标检测 #人工智能 #pytorch

于 2024-03-31 15:09:35 首次发布

文章介绍了如何使用mmengine库中的`get_model_complexity_info`函数来计算深度学习模型的FLOPs（浮点运算次数）和参数数量，适用于单输入和多输入网络，并展示了在不同设备上的计算示例。

使用analysis.get_model_complexity_info()打印任意模型参数和flops

对于只有一个输入的网络

from mmengine.analysis import get_model_complexity_info
from mmengine.analysis.print_helper import _format_size
from models.Models import mynet
import torch

if __name__ == '__main__':

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print('==> Building model..')
    net = mynet()
    net