Function flops

model_zigma.py:1183–1211

(model, shape=(3, 32, 32))

def flops(model, shape=(3, 32, 32)):
    from fvcore.nn import FlopCountAnalysis, flop_count_str, flop_count, parameter_count

    # shape = self.__input_shape__[1:]
    supported_ops = {
        "aten::silu": None,  # as relu is in _IGNORED_OPS
        "aten::neg": None,  # as relu is in _IGNORED_OPS
        "aten::exp": None,  # as relu is in _IGNORED_OPS
        "aten::flip": None,  # as permute is in _IGNORED_OPS
        # "prim::PythonOp.CrossScan": None,
        # "prim::PythonOp.CrossMerge": None,
        "prim::PythonOp.SelectiveScanFn": partial(
            selective_scan_flop_jit, flops_fn=flops_selective_scan_fn
        ),
    }

    model = copy.deepcopy(model)
    model.cuda().eval()

    input = torch.randn((1, *shape), device=next(model.parameters()).device)
    timestep = torch.rand((1), device=next(model.parameters()).device)
    params = parameter_count(model)[""]
    Gflops, unsupported = flop_count(
        model=model, inputs=(input, timestep), supported_ops=supported_ops
    )

    del model, input
    print(f"params {params} GFLOPs {sum(Gflops.values())}")
    return sum(Gflops.values()) * 1e9


if __name__ == "__main__":