使用mmengine.analysis.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 = net.to(device=device)
    input_shape = (3, 256, 256)
    outputs = get_model_complexity_info(net, input_shape=input_shape, show_table=False, show_arch=False)
    flops = _format_size(outputs['flops'])
    params = _format_size(outputs['params'])
    print("flops:{}".format(flops))
    print("params:{}".format(params))

• 对于有两个或多个输入的网络

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 = net.to(device=device)
    input_shape = ((3, 256, 256), (3, 256, 256)) # 如果有更多输入，在元组中再添加即可
    outputs = get_model_complexity_info(net, input_shape=input_shape, show_table=False, show_arch=False)
    flops = _format_size(outputs['flops'])
    params = _format_size(outputs['params'])
    print("flops:{}".format(flops))
    print("params:{}".format(params))

• 输出结果：

==> Building model..
flops:11.525G
params:6.036M