linea/util/profiler.py

import copy
from calflops import calculate_flops
from typing import Tuple

def stats(
    model, args,
    input_shape: Tuple=(1, 3, 640, 640), ) -> Tuple[int, dict]:

    base_size = args.eval_spatial_size[0]
    input_shape = (1, 3, base_size, base_size)

    model_for_info = copy.deepcopy(model).deploy()

    flops, macs, _ = calculate_flops(model=model_for_info,
                                        input_shape=input_shape,
                                        output_as_string=True,
                                        output_precision=4,
                                        print_detailed=False)
    params = sum(p.numel() for p in model_for_info.parameters())
    del model_for_info
    return {'flops': flops, 'macs': macs, 'params': params}