ultralytics-thop 2.0.14__py3-none-any.whl

thop/__init__.py

@@ -0,0 +1,11 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+__version__ = "2.0.14"
+
+
+import torch
+
+from .profile import profile, profile_origin
+from .utils import clever_format
+
+default_dtype = torch.float64

thop/fx_profile.py

@@ -0,0 +1,237 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+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. But You are using an old version PyTorch {torch.__version__}. "
+    )
+
+
+def count_clamp(input_shapes, output_shapes):
+    """Ensures tensor array sizes are appropriate by clamping specified input and output shapes."""
+    return 0
+
+
+def count_mul(input_shapes, output_shapes):
+    """Returns the number of elements in the first output shape."""
+    return output_shapes[0].numel()
+
+
+def count_matmul(input_shapes, output_shapes):
+    """Calculates matrix multiplication ops based on input and output tensor shapes for performance profiling."""
+    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):
+    """Calculates the total FLOPs for a linear layer, including bias operations if specified."""
+    flops = count_matmul(input_shapes, output_shapes)
+    if "bias" in kwargs:
+        flops += output_shapes[0].numel()
+    return flops
+
+
+from .vision.calc_func import calculate_conv
+
+
+def count_fn_conv2d(input_shapes, output_shapes, *args, **kwargs):
+    """Calculates total operations (FLOPs) for a 2D conv layer based on input and output shapes using
+    `calculate_conv`.
+    """
+    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):
+    """Counts the FLOPs for a fully connected (linear) layer in a neural network module."""
+    return count_matmul(input_shapes, output_shapes)
+
+
+def count_zero_ops(module: nn.Module, input_shapes, output_shapes, *args, **kwargs):
+    """Returns 0 for a neural network module, input shapes, and output shapes in PyTorch."""
+    return 0
+
+
+def count_nn_conv2d(module: nn.Conv2d, input_shapes, output_shapes):
+    """Calculates FLOPs for a 2D Conv2D layer in an nn.Module using input and 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):
+    """Calculate FLOPs for an nn.BatchNorm2d layer based on the given output shape."""
+    assert len(output_shapes) == 1, "nn.BatchNorm2d should only have one output"
+    y = output_shapes[0]
+    return 2 * y.numel()
+
+
+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 prRed, prYellow
+
+
+def null_print(*args, **kwargs):
+    """A no-op print function that takes any arguments without performing any actions."""
+    return
+
+
+def fx_profile(mod: nn.Module, input: th.Tensor, verbose=False):
+    """Profiles nn.Module for total FLOPs per operation and prints detailed nodes if verbose."""
+    gm: torch.fx.GraphModule = symbolic_trace(mod)
+    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 = []
+        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 = [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("weight_shape: None")
+            else:
+                print(type(m))
+                print(f"weight_shape: {mod.state_dict()[f'{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):
+            """Performs forward pass on given input data."""
+            return input / 1
+
+    class MyModule(torch.nn.Module):
+        def __init__(self):
+            """Initializes MyModule with two linear layers and a custom MyOP operator."""
+            super().__init__()
+            self.linear1 = torch.nn.Linear(5, 3)
+            self.linear2 = torch.nn.Linear(5, 3)
+            self.myop = MyOP()
+
+        def forward(self, x):
+            """Applies two linear transformations to the input tensor, clamps the second, then combines and processes
+            with MyOP operator.
+            """
+            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/profile.py

@@ -0,0 +1,228 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from thop.rnn_hooks import *
+from thop.vision.basic_hooks import *
+
+from .utils import prRed
+
+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,
+    nn.SyncBatchNorm: count_normalization,
+}
+
+
+def profile_origin(model, inputs, custom_ops=None, verbose=True, report_missing=False):
+    """Profiles a PyTorch model's operations and parameters, applying either custom or default hooks."""
+    handler_collection = []
+    types_collection = set()
+    if custom_ops is None:
+        custom_ops = {}
+    if report_missing:
+        verbose = True
+
+    def add_hooks(m):
+        if list(m.children()):
+            return
+
+        if hasattr(m, "total_ops") or hasattr(m, "total_params"):
+            logging.warning(
+                f"Either .total_ops or .total_params is already defined in {str(m)}. "
+                "Be careful, it might change your code's behavior."
+            )
+
+        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(f"[INFO] Customize rule {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(f"[INFO] Register {fn.__qualname__}() for {m_type}.")
+        else:
+            if m_type not in types_collection and report_missing:
+                prRed(f"[WARN] Cannot find rule for {m_type}. Treat it as zero Macs and zero Params.")
+
+        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 list(m.children()):  # 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 list(m.children()):
+            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,
+):
+    """Profiles a PyTorch model, returning total operations, parameters, and optionally layer-wise details."""
+    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):
+        """Registers hooks to a neural network module to track total operations and parameters."""
+        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(f"[INFO] Customize rule {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(f"[INFO] Register {fn.__qualname__}() for {m_type}.")
+        else:
+            if m_type not in types_collection and report_missing:
+                prRed(f"[WARN] Cannot find rule for {m_type}. Treat it as zero Macs and zero Params.")
+
+        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):
+        """Recursively counts the total operations and parameters of the given PyTorch module and its submodules."""
+        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

thop/rnn_hooks.py

@@ -0,0 +1,202 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import torch
+import torch.nn as nn
+from torch.nn.utils.rnn import PackedSequence
+
+
+def _count_rnn_cell(input_size, hidden_size, bias=True):
+    """Calculate the total operations for an RNN cell given input size, hidden size, and optional bias."""
+    total_ops = hidden_size * (input_size + hidden_size) + hidden_size
+    if bias:
+        total_ops += hidden_size * 2
+
+    return total_ops
+
+
+def count_rnn_cell(m: nn.RNNCell, x: torch.Tensor, y: torch.Tensor):
+    """Counts the total RNN cell operations based on input tensor, hidden size, bias, and batch size."""
+    total_ops = _count_rnn_cell(m.input_size, m.hidden_size, m.bias)
+
+    batch_size = x[0].size(0)
+    total_ops *= batch_size
+
+    m.total_ops += torch.DoubleTensor([int(total_ops)])
+
+
+def _count_gru_cell(input_size, hidden_size, bias=True):
+    """Counts the total operations for a GRU cell based on input size, hidden size, and bias configuration."""
+    total_ops = 0
+    # r = \sigma(W_{ir} x + b_{ir} + W_{hr} h + b_{hr}) \\
+    # z = \sigma(W_{iz} x + b_{iz} + W_{hz} h + b_{hz}) \\
+    state_ops = (hidden_size + input_size) * hidden_size + hidden_size
+    if bias:
+        state_ops += hidden_size * 2
+    total_ops += state_ops * 2
+
+    # n = \tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn})) \\
+    total_ops += (hidden_size + input_size) * hidden_size + hidden_size
+    if bias:
+        total_ops += hidden_size * 2
+    # r hadamard : r * (~)
+    total_ops += hidden_size
+
+    # h' = (1 - z) * n + z * h
+    # hadamard hadamard add
+    total_ops += hidden_size * 3
+
+    return total_ops
+
+
+def count_gru_cell(m: nn.GRUCell, x: torch.Tensor, y: torch.Tensor):
+    """Calculates and updates the total operations for a GRU cell in a mini-batch during inference."""
+    total_ops = _count_gru_cell(m.input_size, m.hidden_size, m.bias)
+
+    batch_size = x[0].size(0)
+    total_ops *= batch_size
+
+    m.total_ops += torch.DoubleTensor([int(total_ops)])
+
+
+def _count_lstm_cell(input_size, hidden_size, bias=True):
+    """Counts LSTM cell operations during inference based on input size, hidden size, and bias configuration."""
+    total_ops = 0
+
+    # i = \sigma(W_{ii} x + b_{ii} + W_{hi} h + b_{hi}) \\
+    # f = \sigma(W_{if} x + b_{if} + W_{hf} h + b_{hf}) \\
+    # o = \sigma(W_{io} x + b_{io} + W_{ho} h + b_{ho}) \\
+    # g = \tanh(W_{ig} x + b_{ig} + W_{hg} h + b_{hg}) \\
+    state_ops = (input_size + hidden_size) * hidden_size + hidden_size
+    if bias:
+        state_ops += hidden_size * 2
+    total_ops += state_ops * 4
+
+    # c' = f * c + i * g \\
+    # hadamard hadamard add
+    total_ops += hidden_size * 3
+
+    # h' = o * \tanh(c') \\
+    total_ops += hidden_size
+
+    return total_ops
+
+
+def count_lstm_cell(m: nn.LSTMCell, x: torch.Tensor, y: torch.Tensor):
+    """Counts and updates the total operations for an LSTM cell in a mini-batch during inference."""
+    total_ops = _count_lstm_cell(m.input_size, m.hidden_size, m.bias)
+
+    batch_size = x[0].size(0)
+    total_ops *= batch_size
+
+    m.total_ops += torch.DoubleTensor([int(total_ops)])
+
+
+def count_rnn(m: nn.RNN, x, y):
+    """Calculate and update the total number of operations for each RNN cell in a given batch."""
+    bias = m.bias
+    input_size = m.input_size
+    hidden_size = m.hidden_size
+    num_layers = m.num_layers
+
+    if isinstance(x[0], PackedSequence):
+        batch_size = torch.max(x[0].batch_sizes)
+        num_steps = x[0].batch_sizes.size(0)
+    elif m.batch_first:
+        batch_size = x[0].size(0)
+        num_steps = x[0].size(1)
+    else:
+        batch_size = x[0].size(1)
+        num_steps = x[0].size(0)
+
+    total_ops = 0
+    if m.bidirectional:
+        total_ops += _count_rnn_cell(input_size, hidden_size, bias) * 2
+    else:
+        total_ops += _count_rnn_cell(input_size, hidden_size, bias)
+
+    for _ in range(num_layers - 1):
+        total_ops += (
+            _count_rnn_cell(hidden_size * 2, hidden_size, bias) * 2
+            if m.bidirectional
+            else _count_rnn_cell(hidden_size, hidden_size, bias)
+        )
+    # time unroll
+    total_ops *= num_steps
+    # batch_size
+    total_ops *= batch_size
+
+    m.total_ops += torch.DoubleTensor([int(total_ops)])
+
+
+def count_gru(m: nn.GRU, x, y):
+    """Calculates total operations for a GRU layer, updating the model's operation count based on batch size."""
+    bias = m.bias
+    input_size = m.input_size
+    hidden_size = m.hidden_size
+    num_layers = m.num_layers
+
+    if isinstance(x[0], PackedSequence):
+        batch_size = torch.max(x[0].batch_sizes)
+        num_steps = x[0].batch_sizes.size(0)
+    elif m.batch_first:
+        batch_size = x[0].size(0)
+        num_steps = x[0].size(1)
+    else:
+        batch_size = x[0].size(1)
+        num_steps = x[0].size(0)
+
+    total_ops = 0
+    if m.bidirectional:
+        total_ops += _count_gru_cell(input_size, hidden_size, bias) * 2
+    else:
+        total_ops += _count_gru_cell(input_size, hidden_size, bias)
+
+    for _ in range(num_layers - 1):
+        total_ops += (
+            _count_gru_cell(hidden_size * 2, hidden_size, bias) * 2
+            if m.bidirectional
+            else _count_gru_cell(hidden_size, hidden_size, bias)
+        )
+    # time unroll
+    total_ops *= num_steps
+    # batch_size
+    total_ops *= batch_size
+
+    m.total_ops += torch.DoubleTensor([int(total_ops)])
+
+
+def count_lstm(m: nn.LSTM, x, y):
+    """Calculate total operations for LSTM layers, including bidirectional, updating model's total operations."""
+    bias = m.bias
+    input_size = m.input_size
+    hidden_size = m.hidden_size
+    num_layers = m.num_layers
+
+    if isinstance(x[0], PackedSequence):
+        batch_size = torch.max(x[0].batch_sizes)
+        num_steps = x[0].batch_sizes.size(0)
+    elif m.batch_first:
+        batch_size = x[0].size(0)
+        num_steps = x[0].size(1)
+    else:
+        batch_size = x[0].size(1)
+        num_steps = x[0].size(0)
+
+    total_ops = 0
+    if m.bidirectional:
+        total_ops += _count_lstm_cell(input_size, hidden_size, bias) * 2
+    else:
+        total_ops += _count_lstm_cell(input_size, hidden_size, bias)
+
+    for _ in range(num_layers - 1):
+        total_ops += (
+            _count_lstm_cell(hidden_size * 2, hidden_size, bias) * 2
+            if m.bidirectional
+            else _count_lstm_cell(hidden_size, hidden_size, bias)
+        )
+    # time unroll
+    total_ops *= num_steps
+    # batch_size
+    total_ops *= batch_size
+
+    m.total_ops += torch.DoubleTensor([int(total_ops)])