ultralytics-thop 0.0.1__py3-none-any.whl

thop/__init__.py

@@ -0,0 +1,8 @@
+# from .onnx_profile import OnnxProfile
+import torch
+
+from .profile import profile, profile_origin
+from .utils import clever_format
+
+default_dtype = torch.float64
+from .__version__ import __version__

thop/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.1.1"

thop/fx_profile.py

@@ -0,0 +1,224 @@
+import logging
+from distutils.version import LooseVersion
+
+import torch
+import torch as th
+import torch.nn as nn
+
+if LooseVersion(torch.__version__) < LooseVersion("1.8.0"):
+    logging.warning(
+        f"torch.fx requires version higher than 1.8.0. "
+        f"But You are using an old version PyTorch {torch.__version__}. "
+    )
+
+
+def count_clamp(input_shapes, output_shapes):
+    return 0
+
+
+def count_mul(input_shapes, output_shapes):
+    # element-wise
+    return output_shapes[0].numel()
+
+
+def count_matmul(input_shapes, output_shapes):
+    in_shape = input_shapes[0]
+    out_shape = output_shapes[0]
+    in_features = in_shape[-1]
+    num_elements = out_shape.numel()
+    return in_features * num_elements
+
+
+def count_fn_linear(input_shapes, output_shapes, *args, **kwargs):
+    mul_flops = count_matmul(input_shapes, output_shapes)
+    if "bias" in kwargs:
+        add_flops = output_shapes[0].numel()
+    return mul_flops
+
+
+from .vision.calc_func import calculate_conv
+
+
+def count_fn_conv2d(input_shapes, output_shapes, *args, **kwargs):
+    inputs, weight, bias, stride, padding, dilation, groups = args
+    if len(input_shapes) == 2:
+        x_shape, k_shape = input_shapes
+    elif len(input_shapes) == 3:
+        x_shape, k_shape, b_shape = input_shapes
+    out_shape = output_shapes[0]
+
+    kernel_parameters = k_shape[2:].numel()
+    bias_op = 0  # check it later
+    in_channel = x_shape[1]
+
+    total_ops = calculate_conv(bias_op, kernel_parameters, out_shape.numel(), in_channel, groups).item()
+    return int(total_ops)
+
+
+def count_nn_linear(module: nn.Module, input_shapes, output_shapes):
+    return count_matmul(input_shapes, output_shapes)
+
+
+def count_zero_ops(module: nn.Module, input_shapes, output_shapes, *args, **kwargs):
+    return 0
+
+
+def count_nn_conv2d(module: nn.Conv2d, input_shapes, output_shapes):
+    bias_op = 1 if module.bias is not None else 0
+    out_shape = output_shapes[0]
+
+    in_channel = module.in_channels
+    groups = module.groups
+    kernel_ops = module.weight.shape[2:].numel()
+    total_ops = calculate_conv(bias_op, kernel_ops, out_shape.numel(), in_channel, groups).item()
+    return int(total_ops)
+
+
+def count_nn_bn2d(module: nn.BatchNorm2d, input_shapes, output_shapes):
+    assert len(output_shapes) == 1, "nn.BatchNorm2d should only have one output"
+    y = output_shapes[0]
+    # y = (x - mean) / \sqrt{var + e} * weight + bias
+    total_ops = 2 * y.numel()
+    return total_ops
+
+
+zero_ops = (
+    nn.ReLU,
+    nn.ReLU6,
+    nn.Dropout,
+    nn.MaxPool2d,
+    nn.AvgPool2d,
+    nn.AdaptiveAvgPool2d,
+)
+
+count_map = {
+    nn.Linear: count_nn_linear,
+    nn.Conv2d: count_nn_conv2d,
+    nn.BatchNorm2d: count_nn_bn2d,
+    "function linear": count_fn_linear,
+    "clamp": count_clamp,
+    "built-in function add": count_zero_ops,
+    "built-in method fl": count_zero_ops,
+    "built-in method conv2d of type object": count_fn_conv2d,
+    "built-in function mul": count_mul,
+    "built-in function truediv": count_mul,
+}
+
+for k in zero_ops:
+    count_map[k] = count_zero_ops
+
+missing_maps = {}
+
+from torch.fx import symbolic_trace
+from torch.fx.passes.shape_prop import ShapeProp
+
+from .utils import prGreen, prRed, prYellow
+
+
+def null_print(*args, **kwargs):
+    return
+
+
+def fx_profile(mod: nn.Module, input: th.Tensor, verbose=False):
+    gm: torch.fx.GraphModule = symbolic_trace(mod)
+    g = gm.graph
+    ShapeProp(gm).propagate(input)
+
+    fprint = null_print
+    if verbose:
+        fprint = print
+
+    v_maps = {}
+    total_flops = 0
+
+    for node in gm.graph.nodes:
+        # print(f"{node.target},\t{node.op},\t{node.meta['tensor_meta'].dtype},\t{node.meta['tensor_meta'].shape}")
+        fprint(f"NodeOP:{node.op},\tTarget:{node.target},\tNodeName:{node.name},\tNodeArgs:{node.args}")
+        # node_op_type = str(node.target).split(".")[-1]
+        node_flops = None
+
+        input_shapes = []
+        output_shapes = []
+        fprint("input_shape:", end="\t")
+        for arg in node.args:
+            if str(arg) not in v_maps:
+                continue
+            fprint(f"{v_maps[str(arg)]}", end="\t")
+            input_shapes.append(v_maps[str(arg)])
+        fprint()
+        fprint(f"output_shape:\t{node.meta['tensor_meta'].shape}")
+        output_shapes.append(node.meta["tensor_meta"].shape)
+
+        if node.op in ["output", "placeholder"]:
+            node_flops = 0
+        elif node.op == "call_function":
+            # torch internal functions
+            key = str(node.target).split("at")[0].replace("<", "").replace(">", "").strip()
+            if key in count_map:
+                node_flops = count_map[key](input_shapes, output_shapes, *node.args, **node.kwargs)
+            else:
+                missing_maps[key] = (node.op, key)
+                prRed(f"|{key}| is missing")
+        elif node.op == "call_method":
+            # torch internal functions
+            # fprint(str(node.target) in count_map, str(node.target), count_map.keys())
+            key = str(node.target)
+            if key in count_map:
+                node_flops = count_map[key](input_shapes, output_shapes)
+            else:
+                missing_maps[key] = (node.op, key)
+                prRed(f"{key} is missing")
+        elif node.op == "call_module":
+            # torch.nn modules
+            # m = getattr(mod, node.target, None)
+            m = mod.get_submodule(node.target)
+            key = type(m)
+            fprint(type(m), type(m) in count_map)
+            if type(m) in count_map:
+                node_flops = count_map[type(m)](m, input_shapes, output_shapes)
+            else:
+                missing_maps[key] = (node.op,)
+                prRed(f"{key} is missing")
+            print("module type:", type(m))
+            if isinstance(m, zero_ops):
+                print(f"weight_shape: None")
+            else:
+                print(type(m))
+                print(f"weight_shape: {mod.state_dict()[node.target + '.weight'].shape}")
+
+        v_maps[str(node.name)] = node.meta["tensor_meta"].shape
+        if node_flops is not None:
+            total_flops += node_flops
+        prYellow(f"Current node's FLOPs: {node_flops}, total FLOPs: {total_flops}")
+        fprint("==" * 40)
+
+    if len(missing_maps.keys()) > 0:
+        from pprint import pprint
+
+        print("Missing operators: ")
+        pprint(missing_maps)
+    return total_flops
+
+
+if __name__ == "__main__":
+
+    class MyOP(nn.Module):
+        def forward(self, input):
+            return input / 1
+
+    class MyModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear1 = torch.nn.Linear(5, 3)
+            self.linear2 = torch.nn.Linear(5, 3)
+            self.myop = MyOP()
+
+        def forward(self, x):
+            out1 = self.linear1(x)
+            out2 = self.linear2(x).clamp(min=0.0, max=1.0)
+            return self.myop(out1 + out2)
+
+    net = MyModule()
+    data = th.randn(20, 5)
+    flops = fx_profile(net, data, verbose=False)
+    print(flops)

thop/onnx_profile.py

@@ -0,0 +1,76 @@
+import numpy as np
+import onnx
+import torch
+import torch.nn
+from onnx import numpy_helper
+
+from thop.vision.onnx_counter import onnx_operators
+
+
+class OnnxProfile:
+    def __init__(self):
+        pass
+
+    def calculate_params(self, model: onnx.ModelProto):
+        onnx_weights = model.graph.initializer
+        params = 0
+
+        for onnx_w in onnx_weights:
+            try:
+                weight = numpy_helper.to_array(onnx_w)
+                params += np.prod(weight.shape)
+            except Exception as _:
+                pass
+
+        return params
+
+    def create_dict(self, weight, input, output):
+        diction = {}
+        for w in weight:
+            dim = np.array(w.dims)
+            diction[str(w.name)] = dim
+            if dim.size == 1:
+                diction[str(w.name)] = np.append(1, dim)
+        for i in input:
+            # print(i.type.tensor_type.shape.dim[0].dim_value)
+            dim = np.array(i.type.tensor_type.shape.dim[0].dim_value)
+            # print(i.type.tensor_type.shape.dim.__sizeof__())
+            # name2dims[str(i.name)] = [dim]
+            dim = []
+            for key in i.type.tensor_type.shape.dim:
+                dim = np.append(dim, int(key.dim_value))
+                # print(key.dim_value)
+            # print(dim)
+            diction[str(i.name)] = dim
+            if dim.size == 1:
+                diction[str(i.name)] = np.append(1, dim)
+        for o in output:
+            dim = np.array(o.type.tensor_type.shape.dim[0].dim_value)
+            diction[str(o.name)] = [dim]
+            if dim.size == 1:
+                diction[str(o.name)] = np.append(1, dim)
+        return diction
+
+    def nodes_counter(self, diction, node):
+        if node.op_type not in onnx_operators:
+            print("Sorry, we haven't add ", node.op_type, "into dictionary.")
+            return 0, None, None
+        else:
+            fn = onnx_operators[node.op_type]
+            return fn(diction, node)
+
+    def calculate_macs(self, model: onnx.ModelProto) -> torch.DoubleTensor:
+        macs = 0
+        name2dims = {}
+        weight = model.graph.initializer
+        nodes = model.graph.node
+        input = model.graph.input
+        output = model.graph.output
+        name2dims = self.create_dict(weight, input, output)
+        macs = 0
+        for n in nodes:
+            macs_adding, out_size, outname = self.nodes_counter(name2dims, n)
+
+            name2dims[outname] = out_size
+            macs += macs_adding
+        return np.array(macs[0])

thop/profile.py

@@ -0,0 +1,233 @@
+from distutils.version import LooseVersion
+
+from thop.rnn_hooks import *
+from thop.vision.basic_hooks import *
+
+# logger = logging.getLogger(__name__)
+# logger.setLevel(logging.INFO)
+from .utils import prGreen, prRed, prYellow
+
+if LooseVersion(torch.__version__) < LooseVersion("1.0.0"):
+    logging.warning(
+        "You are using an old version PyTorch {version}, which THOP does NOT support.".format(version=torch.__version__)
+    )
+
+default_dtype = torch.float64
+
+register_hooks = {
+    nn.ZeroPad2d: zero_ops,  # padding does not involve any multiplication.
+    nn.Conv1d: count_convNd,
+    nn.Conv2d: count_convNd,
+    nn.Conv3d: count_convNd,
+    nn.ConvTranspose1d: count_convNd,
+    nn.ConvTranspose2d: count_convNd,
+    nn.ConvTranspose3d: count_convNd,
+    nn.BatchNorm1d: count_normalization,
+    nn.BatchNorm2d: count_normalization,
+    nn.BatchNorm3d: count_normalization,
+    nn.LayerNorm: count_normalization,
+    nn.InstanceNorm1d: count_normalization,
+    nn.InstanceNorm2d: count_normalization,
+    nn.InstanceNorm3d: count_normalization,
+    nn.PReLU: count_prelu,
+    nn.Softmax: count_softmax,
+    nn.ReLU: zero_ops,
+    nn.ReLU6: zero_ops,
+    nn.LeakyReLU: count_relu,
+    nn.MaxPool1d: zero_ops,
+    nn.MaxPool2d: zero_ops,
+    nn.MaxPool3d: zero_ops,
+    nn.AdaptiveMaxPool1d: zero_ops,
+    nn.AdaptiveMaxPool2d: zero_ops,
+    nn.AdaptiveMaxPool3d: zero_ops,
+    nn.AvgPool1d: count_avgpool,
+    nn.AvgPool2d: count_avgpool,
+    nn.AvgPool3d: count_avgpool,
+    nn.AdaptiveAvgPool1d: count_adap_avgpool,
+    nn.AdaptiveAvgPool2d: count_adap_avgpool,
+    nn.AdaptiveAvgPool3d: count_adap_avgpool,
+    nn.Linear: count_linear,
+    nn.Dropout: zero_ops,
+    nn.Upsample: count_upsample,
+    nn.UpsamplingBilinear2d: count_upsample,
+    nn.UpsamplingNearest2d: count_upsample,
+    nn.RNNCell: count_rnn_cell,
+    nn.GRUCell: count_gru_cell,
+    nn.LSTMCell: count_lstm_cell,
+    nn.RNN: count_rnn,
+    nn.GRU: count_gru,
+    nn.LSTM: count_lstm,
+    nn.Sequential: zero_ops,
+    nn.PixelShuffle: zero_ops,
+}
+
+if LooseVersion(torch.__version__) >= LooseVersion("1.1.0"):
+    register_hooks.update({nn.SyncBatchNorm: count_normalization})
+
+
+def profile_origin(model, inputs, custom_ops=None, verbose=True, report_missing=False):
+    handler_collection = []
+    types_collection = set()
+    if custom_ops is None:
+        custom_ops = {}
+    if report_missing:
+        verbose = True
+
+    def add_hooks(m):
+        if len(list(m.children())) > 0:
+            return
+
+        if hasattr(m, "total_ops") or hasattr(m, "total_params"):
+            logging.warning(
+                "Either .total_ops or .total_params is already defined in %s. "
+                "Be careful, it might change your code's behavior." % str(m)
+            )
+
+        m.register_buffer("total_ops", torch.zeros(1, dtype=default_dtype))
+        m.register_buffer("total_params", torch.zeros(1, dtype=default_dtype))
+
+        for p in m.parameters():
+            m.total_params += torch.DoubleTensor([p.numel()])
+
+        m_type = type(m)
+
+        fn = None
+        if m_type in custom_ops:  # if defined both op maps, use custom_ops to overwrite.
+            fn = custom_ops[m_type]
+            if m_type not in types_collection and verbose:
+                print("[INFO] Customize rule %s() %s." % (fn.__qualname__, m_type))
+        elif m_type in register_hooks:
+            fn = register_hooks[m_type]
+            if m_type not in types_collection and verbose:
+                print("[INFO] Register %s() for %s." % (fn.__qualname__, m_type))
+        else:
+            if m_type not in types_collection and report_missing:
+                prRed("[WARN] Cannot find rule for %s. Treat it as zero Macs and zero Params." % m_type)
+
+        if fn is not None:
+            handler = m.register_forward_hook(fn)
+            handler_collection.append(handler)
+        types_collection.add(m_type)
+
+    training = model.training
+
+    model.eval()
+    model.apply(add_hooks)
+
+    with torch.no_grad():
+        model(*inputs)
+
+    total_ops = 0
+    total_params = 0
+    for m in model.modules():
+        if len(list(m.children())) > 0:  # skip for non-leaf module
+            continue
+        total_ops += m.total_ops
+        total_params += m.total_params
+
+    total_ops = total_ops.item()
+    total_params = total_params.item()
+
+    # reset model to original status
+    model.train(training)
+    for handler in handler_collection:
+        handler.remove()
+
+    # remove temporal buffers
+    for n, m in model.named_modules():
+        if len(list(m.children())) > 0:
+            continue
+        if "total_ops" in m._buffers:
+            m._buffers.pop("total_ops")
+        if "total_params" in m._buffers:
+            m._buffers.pop("total_params")
+
+    return total_ops, total_params
+
+
+def profile(
+    model: nn.Module,
+    inputs,
+    custom_ops=None,
+    verbose=True,
+    ret_layer_info=False,
+    report_missing=False,
+):
+    handler_collection = {}
+    types_collection = set()
+    if custom_ops is None:
+        custom_ops = {}
+    if report_missing:
+        # overwrite `verbose` option when enable report_missing
+        verbose = True
+
+    def add_hooks(m: nn.Module):
+        m.register_buffer("total_ops", torch.zeros(1, dtype=torch.float64))
+        m.register_buffer("total_params", torch.zeros(1, dtype=torch.float64))
+
+        # for p in m.parameters():
+        #     m.total_params += torch.DoubleTensor([p.numel()])
+
+        m_type = type(m)
+
+        fn = None
+        if m_type in custom_ops:
+            # if defined both op maps, use custom_ops to overwrite.
+            fn = custom_ops[m_type]
+            if m_type not in types_collection and verbose:
+                print("[INFO] Customize rule %s() %s." % (fn.__qualname__, m_type))
+        elif m_type in register_hooks:
+            fn = register_hooks[m_type]
+            if m_type not in types_collection and verbose:
+                print("[INFO] Register %s() for %s." % (fn.__qualname__, m_type))
+        else:
+            if m_type not in types_collection and report_missing:
+                prRed("[WARN] Cannot find rule for %s. Treat it as zero Macs and zero Params." % m_type)
+
+        if fn is not None:
+            handler_collection[m] = (
+                m.register_forward_hook(fn),
+                m.register_forward_hook(count_parameters),
+            )
+        types_collection.add(m_type)
+
+    prev_training_status = model.training
+
+    model.eval()
+    model.apply(add_hooks)
+
+    with torch.no_grad():
+        model(*inputs)
+
+    def dfs_count(module: nn.Module, prefix="\t") -> (int, int):
+        total_ops, total_params = module.total_ops.item(), 0
+        ret_dict = {}
+        for n, m in module.named_children():
+            # if not hasattr(m, "total_ops") and not hasattr(m, "total_params"):  # and len(list(m.children())) > 0:
+            #     m_ops, m_params = dfs_count(m, prefix=prefix + "\t")
+            # else:
+            #     m_ops, m_params = m.total_ops, m.total_params
+            next_dict = {}
+            if m in handler_collection and not isinstance(m, (nn.Sequential, nn.ModuleList)):
+                m_ops, m_params = m.total_ops.item(), m.total_params.item()
+            else:
+                m_ops, m_params, next_dict = dfs_count(m, prefix=prefix + "\t")
+            ret_dict[n] = (m_ops, m_params, next_dict)
+            total_ops += m_ops
+            total_params += m_params
+        # print(prefix, module._get_name(), (total_ops, total_params))
+        return total_ops, total_params, ret_dict
+
+    total_ops, total_params, ret_dict = dfs_count(model)
+
+    # reset model to original status
+    model.train(prev_training_status)
+    for m, (op_handler, params_handler) in handler_collection.items():
+        op_handler.remove()
+        params_handler.remove()
+        m._buffers.pop("total_ops")
+        m._buffers.pop("total_params")
+
+    if ret_layer_info:
+        return total_ops, total_params, ret_dict
+    return total_ops, total_params