From 1b4439ed63530c94538cb396dfb7f4191e6631f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vincent=20L=C3=A9b=C3=A9?= Date: Tue, 25 Nov 2025 15:16:54 +0100 Subject: [PATCH 1/6] feat: Updating hook to support multiple types for inputs and outputs --- torchsummary/torchsummary.py | 31 +++++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/torchsummary/torchsummary.py b/torchsummary/torchsummary.py index 1ed065f..0a77dd0 100644 --- a/torchsummary/torchsummary.py +++ b/torchsummary/torchsummary.py @@ -5,6 +5,22 @@ from collections import OrderedDict import numpy as np +def extract_shapes(x): + """Extract all tensor shapes from nested structures.""" + shapes = [] + + if torch.is_tensor(x): + shapes.append(list(x.size())) + elif isinstance(x, dict): + for v in x.values(): + shapes.extend(extract_shapes(v)) + elif isinstance(x, (list, tuple)): + for v in x: + shapes.extend(extract_shapes(v)) + # ignore other types (ints, strings, None, etc.) + + return shapes + def summary(model, input_size, batch_size=-1, device=torch.device('cuda:0'), dtypes=None): result, params_info = summary_string( @@ -27,15 +43,8 @@ def hook(module, input, output): m_key = "%s-%i" % (class_name, module_idx + 1) summary[m_key] = OrderedDict() - summary[m_key]["input_shape"] = list(input[0].size()) - summary[m_key]["input_shape"][0] = batch_size - if isinstance(output, (list, tuple)): - summary[m_key]["output_shape"] = [ - [-1] + list(o.size())[1:] for o in output - ] - else: - summary[m_key]["output_shape"] = list(output.size()) - summary[m_key]["output_shape"][0] = batch_size + summary[m_key]["input_shape"] = extract_shapes(input) + summary[m_key]["output_shape"] = extract_shapes(output) params = 0 if hasattr(module, "weight") and hasattr(module.weight, "size"): @@ -91,7 +100,9 @@ def hook(module, input, output): ) total_params += summary[layer]["nb_params"] - total_output += np.prod(summary[layer]["output_shape"]) + for shape in summary[layer]["output_shape"]: + total_output += np.prod(shape) + if "trainable" in summary[layer]: if summary[layer]["trainable"] == True: trainable_params += summary[layer]["nb_params"] From f2a4516bc648429566aafce3cdaf7e6b649df997 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vincent=20L=C3=A9b=C3=A9?= Date: Tue, 25 Nov 2025 17:30:40 +0100 Subject: [PATCH 2/6] fix: Fixing display for layers with multiple outputs --- torchsummary/torchsummary.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/torchsummary/torchsummary.py b/torchsummary/torchsummary.py index 0a77dd0..534bbfc 100644 --- a/torchsummary/torchsummary.py +++ b/torchsummary/torchsummary.py @@ -21,6 +21,17 @@ def extract_shapes(x): return shapes +def format_shapes(layer_name, shapes, nb_params, width): + lines = [] + for i, s in enumerate(shapes): + if i == 0: + line = "{:>{w1}} {:>25} {:>15}".format(layer_name, '[' + str(s) + (',' if len(shapes) > 1 else ']'), nb_params, w1=width) + elif i == len(shapes) - 1: + line = "{:>{w1}} {:>25} {:>15}".format('', str(s) + ']', '', w1=width) + else: + line = "{:>{w1}} {:>25} {:>15}".format('', str(s) + ',', '', w1=width) + lines.append(line) + return "\n".join(lines) def summary(model, input_size, batch_size=-1, device=torch.device('cuda:0'), dtypes=None): result, params_info = summary_string( @@ -93,11 +104,9 @@ def hook(module, input, output): trainable_params = 0 for layer in summary: # input_shape, output_shape, trainable, nb_params - line_new = "{:>20} {:>25} {:>15}".format( - layer, - str(summary[layer]["output_shape"]), - "{0:,}".format(summary[layer]["nb_params"]), - ) + shapes = summary[layer]["output_shape"] + line_new = format_shapes(layer, shapes, summary[layer]["nb_params"], max_layer_len) + total_params += summary[layer]["nb_params"] for shape in summary[layer]["output_shape"]: From 0f09ee75c90e2df4720828b1de4ab9096081d3a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vincent=20L=C3=A9b=C3=A9?= Date: Tue, 25 Nov 2025 17:31:10 +0100 Subject: [PATCH 3/6] feat: Dynamic display based on layers name --- torchsummary/torchsummary.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/torchsummary/torchsummary.py b/torchsummary/torchsummary.py index 534bbfc..902f2f5 100644 --- a/torchsummary/torchsummary.py +++ b/torchsummary/torchsummary.py @@ -94,11 +94,13 @@ def hook(module, input, output): for h in hooks: h.remove() - summary_str += "----------------------------------------------------------------" + "\n" - line_new = "{:>20} {:>25} {:>15}".format( - "Layer (type)", "Output Shape", "Param #") + max_layer_len = max(len(layer) for layer in summary) + + summary_str += "-" * (max_layer_len + 43) + "\n" + line_new = "{:>{w1}} {:>25} {:>15}".format( + "Layer (type)", "Output Shape", "Param #", w1=max_layer_len) summary_str += line_new + "\n" - summary_str += "================================================================" + "\n" + summary_str += "=" * (max_layer_len + 43) + "\n" total_params = 0 total_output = 0 trainable_params = 0 @@ -125,16 +127,16 @@ def hook(module, input, output): total_params_size = abs(total_params * 4. / (1024 ** 2.)) total_size = total_params_size + total_output_size + total_input_size - summary_str += "================================================================" + "\n" + summary_str += "="*(max_layer_len + 43) + "\n" summary_str += "Total params: {0:,}".format(total_params) + "\n" summary_str += "Trainable params: {0:,}".format(trainable_params) + "\n" summary_str += "Non-trainable params: {0:,}".format(total_params - trainable_params) + "\n" - summary_str += "----------------------------------------------------------------" + "\n" + summary_str += "-"*(max_layer_len + 43) + "\n" summary_str += "Input size (MB): %0.2f" % total_input_size + "\n" summary_str += "Forward/backward pass size (MB): %0.2f" % total_output_size + "\n" summary_str += "Params size (MB): %0.2f" % total_params_size + "\n" summary_str += "Estimated Total Size (MB): %0.2f" % total_size + "\n" - summary_str += "----------------------------------------------------------------" + "\n" + summary_str += "-"*(max_layer_len + 43) + "\n" # return summary return summary_str, (total_params, trainable_params) From 2914d0d5d32c9ab2035b7621c5caff5a4f8531d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vincent=20L=C3=A9b=C3=A9?= Date: Tue, 25 Nov 2025 17:50:27 +0100 Subject: [PATCH 4/6] chore: Formating file --- torchsummary/torchsummary.py | 100 ++++++++++++++++++++++++----------- 1 file changed, 68 insertions(+), 32 deletions(-) diff --git a/torchsummary/torchsummary.py b/torchsummary/torchsummary.py index 902f2f5..49186fa 100644 --- a/torchsummary/torchsummary.py +++ b/torchsummary/torchsummary.py @@ -1,12 +1,21 @@ +from collections import OrderedDict +from typing import Any, List + +import numpy as np import torch import torch.nn as nn from torch.autograd import Variable -from collections import OrderedDict -import numpy as np -def extract_shapes(x): - """Extract all tensor shapes from nested structures.""" +def extract_shapes(x: Any) -> List[List[int]]: + """Extract all tensor shapes from nested structures. + + Args: + x (Any): Input which can be a tensor, list, tuple, or dict. + + Returns: + List[List[int]]: A list of shapes, where each shape is represented as a list of integers. + """ shapes = [] if torch.is_tensor(x): @@ -17,35 +26,57 @@ def extract_shapes(x): elif isinstance(x, (list, tuple)): for v in x: shapes.extend(extract_shapes(v)) - # ignore other types (ints, strings, None, etc.) return shapes -def format_shapes(layer_name, shapes, nb_params, width): + +def format_shapes( + layer_name: str, shapes: list, nb_params: torch.Tensor, width: int +) -> str: + """Generate formatted string for layer shapes. + + Args: + layer_name (str): _name of the layer + shapes (list): _list of shapes_ + nb_params (torch.Tensor): _number of parameters_ + width (int): _maximum width of layer names for dynamic formatting_ + + Returns: + str: Formatted string representing the layer, its shapes and number of parameters. + """ lines = [] for i, s in enumerate(shapes): if i == 0: - line = "{:>{w1}} {:>25} {:>15}".format(layer_name, '[' + str(s) + (',' if len(shapes) > 1 else ']'), nb_params, w1=width) + line = "{:>{w1}} {:>25} {:>15}".format( + layer_name, + "[" + str(s) + ("," if len(shapes) > 1 else "]"), + nb_params, + w1=width, + ) elif i == len(shapes) - 1: - line = "{:>{w1}} {:>25} {:>15}".format('', str(s) + ']', '', w1=width) + line = "{:>{w1}} {:>25} {:>15}".format("", str(s) + "]", "", w1=width) else: - line = "{:>{w1}} {:>25} {:>15}".format('', str(s) + ',', '', w1=width) + line = "{:>{w1}} {:>25} {:>15}".format("", str(s) + ",", "", w1=width) lines.append(line) return "\n".join(lines) -def summary(model, input_size, batch_size=-1, device=torch.device('cuda:0'), dtypes=None): - result, params_info = summary_string( - model, input_size, batch_size, device, dtypes) + +def summary( + model, input_size, batch_size=-1, device=torch.device("cuda:0"), dtypes=None +): + result, params_info = summary_string(model, input_size, batch_size, device, dtypes) print(result) return params_info -def summary_string(model, input_size, batch_size=-1, device=torch.device('cuda:0'), dtypes=None): +def summary_string( + model, input_size, batch_size=-1, device=torch.device("cuda:0"), dtypes=None +): if dtypes == None: - dtypes = [torch.FloatTensor]*len(input_size) + dtypes = [torch.FloatTensor] * len(input_size) - summary_str = '' + summary_str = "" def register_hook(module): def hook(module, input, output): @@ -65,9 +96,8 @@ def hook(module, input, output): params += torch.prod(torch.LongTensor(list(module.bias.size()))) summary[m_key]["nb_params"] = params - if ( - not isinstance(module, nn.Sequential) - and not isinstance(module, nn.ModuleList) + if not isinstance(module, nn.Sequential) and not isinstance( + module, nn.ModuleList ): hooks.append(module.register_forward_hook(hook)) @@ -76,8 +106,10 @@ def hook(module, input, output): input_size = [input_size] # batch_size of 2 for batchnorm - x = [torch.rand(2, *in_size).type(dtype).to(device=device) - for in_size, dtype in zip(input_size, dtypes)] + x = [ + torch.rand(2, *in_size).type(dtype).to(device=device) + for in_size, dtype in zip(input_size, dtypes) + ] # create properties summary = OrderedDict() @@ -98,7 +130,8 @@ def hook(module, input, output): summary_str += "-" * (max_layer_len + 43) + "\n" line_new = "{:>{w1}} {:>25} {:>15}".format( - "Layer (type)", "Output Shape", "Param #", w1=max_layer_len) + "Layer (type)", "Output Shape", "Param #", w1=max_layer_len + ) summary_str += line_new + "\n" summary_str += "=" * (max_layer_len + 43) + "\n" total_params = 0 @@ -107,7 +140,9 @@ def hook(module, input, output): for layer in summary: # input_shape, output_shape, trainable, nb_params shapes = summary[layer]["output_shape"] - line_new = format_shapes(layer, shapes, summary[layer]["nb_params"], max_layer_len) + line_new = format_shapes( + layer, shapes, summary[layer]["nb_params"], max_layer_len + ) total_params += summary[layer]["nb_params"] @@ -120,23 +155,24 @@ def hook(module, input, output): summary_str += line_new + "\n" # assume 4 bytes/number (float on cuda). - total_input_size = abs(np.prod(sum(input_size, ())) - * batch_size * 4. / (1024 ** 2.)) - total_output_size = abs(2. * total_output * 4. / - (1024 ** 2.)) # x2 for gradients - total_params_size = abs(total_params * 4. / (1024 ** 2.)) + total_input_size = abs( + np.prod(sum(input_size, ())) * batch_size * 4.0 / (1024**2.0) + ) + total_output_size = abs(2.0 * total_output * 4.0 / (1024**2.0)) # x2 for gradients + total_params_size = abs(total_params * 4.0 / (1024**2.0)) total_size = total_params_size + total_output_size + total_input_size - summary_str += "="*(max_layer_len + 43) + "\n" + summary_str += "=" * (max_layer_len + 43) + "\n" summary_str += "Total params: {0:,}".format(total_params) + "\n" summary_str += "Trainable params: {0:,}".format(trainable_params) + "\n" - summary_str += "Non-trainable params: {0:,}".format(total_params - - trainable_params) + "\n" - summary_str += "-"*(max_layer_len + 43) + "\n" + summary_str += ( + "Non-trainable params: {0:,}".format(total_params - trainable_params) + "\n" + ) + summary_str += "-" * (max_layer_len + 43) + "\n" summary_str += "Input size (MB): %0.2f" % total_input_size + "\n" summary_str += "Forward/backward pass size (MB): %0.2f" % total_output_size + "\n" summary_str += "Params size (MB): %0.2f" % total_params_size + "\n" summary_str += "Estimated Total Size (MB): %0.2f" % total_size + "\n" - summary_str += "-"*(max_layer_len + 43) + "\n" + summary_str += "-" * (max_layer_len + 43) + "\n" # return summary return summary_str, (total_params, trainable_params) From e74448b27bb4afdae7f226645007d13f5f73f74c Mon Sep 17 00:00:00 2001 From: Vincent Lebe Date: Wed, 25 Mar 2026 14:12:32 +0100 Subject: [PATCH 5/6] feat: Ignore parametrized layers --- torchsummary/torchsummary.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/torchsummary/torchsummary.py b/torchsummary/torchsummary.py index 49186fa..682e399 100644 --- a/torchsummary/torchsummary.py +++ b/torchsummary/torchsummary.py @@ -78,7 +78,7 @@ def summary_string( summary_str = "" - def register_hook(module): + def register_hook(module_name, module): def hook(module, input, output): class_name = str(module.__class__).split(".")[-1].split("'")[0] module_idx = len(summary) @@ -96,8 +96,11 @@ def hook(module, input, output): params += torch.prod(torch.LongTensor(list(module.bias.size()))) summary[m_key]["nb_params"] = params - if not isinstance(module, nn.Sequential) and not isinstance( - module, nn.ModuleList + name_parts = module_name.split(".") if module_name else [] + if ( + not isinstance(module, nn.Sequential) + and not isinstance(module, nn.ModuleList) + and "parametrizations" not in name_parts ): hooks.append(module.register_forward_hook(hook)) @@ -116,7 +119,8 @@ def hook(module, input, output): hooks = [] # register hook - model.apply(register_hook) + for module_name, module in model.named_modules(): + register_hook(module_name, module) # make a forward pass # print(x.shape) From 8354aea835f3d638e70d56aed0ec448a7b5ce565 Mon Sep 17 00:00:00 2001 From: Vincent Lebe Date: Tue, 21 Jul 2026 14:17:17 +0200 Subject: [PATCH 6/6] fix: update parametrized layers display --- .../tests/unit_tests/torchsummary_test.py | 90 +++++++++++++++++-- torchsummary/torchsummary.py | 67 ++++++++------ 2 files changed, 123 insertions(+), 34 deletions(-) diff --git a/torchsummary/tests/unit_tests/torchsummary_test.py b/torchsummary/tests/unit_tests/torchsummary_test.py index ec4b33e..2d9c09b 100644 --- a/torchsummary/tests/unit_tests/torchsummary_test.py +++ b/torchsummary/tests/unit_tests/torchsummary_test.py @@ -1,10 +1,18 @@ import unittest -from torchsummary import summary, summary_string -from torchsummary.tests.test_models.test_model import SingleInputNet, MultipleInputNet, MultipleInputNetDifferentDtypes + import torch +import torch.nn.utils.parametrize as parametrize + +from torchsummary import summary, summary_string +from torchsummary.tests.test_models.test_model import ( + MultipleInputNet, + MultipleInputNetDifferentDtypes, + SingleInputNet, +) gpu_if_available = "cuda:0" if torch.cuda.is_available() else "cpu" + class torchsummaryTests(unittest.TestCase): def test_single_input(self): model = SingleInputNet() @@ -17,8 +25,7 @@ def test_multiple_input(self): model = MultipleInputNet() input1 = (1, 300) input2 = (1, 300) - total_params, trainable_params = summary( - model, [input1, input2], device="cpu") + total_params, trainable_params = summary(model, [input1, input2], device="cpu") self.assertEqual(total_params, 31120) self.assertEqual(trainable_params, 31120) @@ -44,7 +51,8 @@ def test_multiple_input_types(self): input2 = (1, 300) dtypes = [torch.FloatTensor, torch.LongTensor] total_params, trainable_params = summary( - model, [input1, input2], device="cpu", dtypes=dtypes) + model, [input1, input2], device="cpu", dtypes=dtypes + ) self.assertEqual(total_params, 31120) self.assertEqual(trainable_params, 31120) @@ -54,11 +62,79 @@ def test_single_input(self): model = SingleInputNet() input = (1, 28, 28) result, (total_params, trainable_params) = summary_string( - model, input, device="cpu") + model, input, device="cpu" + ) self.assertEqual(type(result), str) self.assertEqual(total_params, 21840) self.assertEqual(trainable_params, 21840) + def test_frozen_bias_accounting(self): + model = torch.nn.Linear(2, 5) + model.bias.requires_grad = False + _, (total_params, trainable_params) = summary_string( + model, (1, 2), device="cpu" + ) + self.assertEqual(total_params, 15) + self.assertEqual(trainable_params, 10) + + def test_non_weight_parameter_is_counted(self): + class ModelWithExtraParameter(torch.nn.Module): + def __init__(self): + super(ModelWithExtraParameter, self).__init__() + self.linear = torch.nn.Linear(2, 2, bias=False) + self.extra = torch.nn.Parameter(torch.ones(3)) + + def forward(self, x): + return self.linear(x) * self.extra[:2] + + model = ModelWithExtraParameter() + _, (total_params, trainable_params) = summary_string( + model, (1, 2), device="cpu" + ) + self.assertEqual(total_params, 7) + self.assertEqual(trainable_params, 7) + + def test_parametrization_submodules_are_filtered(self): + class IdentityParametrization(torch.nn.Module): + def forward(self, x): + return x + + model = torch.nn.Linear(2, 2) + parametrize.register_parametrization(model, "weight", IdentityParametrization()) + result, (total_params, trainable_params) = summary_string( + model, (1, 2), device="cpu" + ) + + self.assertEqual(total_params, 6) + self.assertEqual(trainable_params, 6) + self.assertNotIn("IdentityParametrization", result) + self.assertNotIn("Displayed (effective) params", result) + self.assertNotIn("Buffer params", result) + + def test_parametrization_owned_parameters_are_counted_on_parent_layer(self): + class ParametrizationWithAuxParameter(torch.nn.Module): + def __init__(self): + super(ParametrizationWithAuxParameter, self).__init__() + self.aux = torch.nn.Parameter(torch.ones(3), requires_grad=False) + + def forward(self, x): + return x + + model = torch.nn.Linear(2, 2) + parametrize.register_parametrization( + model, "weight", ParametrizationWithAuxParameter() + ) + result, (total_params, trainable_params) = summary_string( + model, (1, 2), device="cpu" + ) + + self.assertEqual(total_params, 9) + self.assertEqual(trainable_params, 6) + param_line = next( + line for line in result.splitlines() if "ParametrizedLinear-1" in line + ) + self.assertTrue(param_line.rstrip().endswith("9")) + -if __name__ == '__main__': +if __name__ == "__main__": unittest.main(buffer=True) diff --git a/torchsummary/torchsummary.py b/torchsummary/torchsummary.py index 682e399..6f26dc6 100644 --- a/torchsummary/torchsummary.py +++ b/torchsummary/torchsummary.py @@ -4,7 +4,6 @@ import numpy as np import torch import torch.nn as nn -from torch.autograd import Variable def extract_shapes(x: Any) -> List[List[int]]: @@ -61,6 +60,35 @@ def format_shapes( return "\n".join(lines) +def should_register_hook(module_name: str, module: nn.Module) -> bool: + name_parts = module_name.split(".") if module_name else [] + return ( + not isinstance(module, nn.Sequential) + and not isinstance(module, nn.ModuleList) + and "parametrizations" not in name_parts + ) + + +def parameter_owner_name(parameter_name: str) -> str: + if parameter_name.startswith("parametrizations."): + return "" + if ".parametrizations." in parameter_name: + return parameter_name.split(".parametrizations.", 1)[0] + if "." in parameter_name: + return parameter_name.rsplit(".", 1)[0] + return "" + + +def get_module_param_counts(model: nn.Module) -> dict: + module_param_counts = {} + for parameter_name, parameter in model.named_parameters(): + owner_name = parameter_owner_name(parameter_name) + module_param_counts[owner_name] = ( + module_param_counts.get(owner_name, 0) + int(parameter.numel()) + ) + return module_param_counts + + def summary( model, input_size, batch_size=-1, device=torch.device("cuda:0"), dtypes=None ): @@ -73,10 +101,11 @@ def summary( def summary_string( model, input_size, batch_size=-1, device=torch.device("cuda:0"), dtypes=None ): - if dtypes == None: + if dtypes is None: dtypes = [torch.FloatTensor] * len(input_size) summary_str = "" + module_param_counts = get_module_param_counts(model) def register_hook(module_name, module): def hook(module, input, output): @@ -88,20 +117,9 @@ def hook(module, input, output): summary[m_key]["input_shape"] = extract_shapes(input) summary[m_key]["output_shape"] = extract_shapes(output) - params = 0 - if hasattr(module, "weight") and hasattr(module.weight, "size"): - params += torch.prod(torch.LongTensor(list(module.weight.size()))) - summary[m_key]["trainable"] = module.weight.requires_grad - if hasattr(module, "bias") and hasattr(module.bias, "size"): - params += torch.prod(torch.LongTensor(list(module.bias.size()))) - summary[m_key]["nb_params"] = params - - name_parts = module_name.split(".") if module_name else [] - if ( - not isinstance(module, nn.Sequential) - and not isinstance(module, nn.ModuleList) - and "parametrizations" not in name_parts - ): + summary[m_key]["nb_params"] = module_param_counts.get(module_name, 0) + + if should_register_hook(module_name, module): hooks.append(module.register_forward_hook(hook)) # multiple inputs to the network @@ -138,9 +156,7 @@ def hook(module, input, output): ) summary_str += line_new + "\n" summary_str += "=" * (max_layer_len + 43) + "\n" - total_params = 0 total_output = 0 - trainable_params = 0 for layer in summary: # input_shape, output_shape, trainable, nb_params shapes = summary[layer]["output_shape"] @@ -148,16 +164,15 @@ def hook(module, input, output): layer, shapes, summary[layer]["nb_params"], max_layer_len ) - total_params += summary[layer]["nb_params"] - for shape in summary[layer]["output_shape"]: total_output += np.prod(shape) - - if "trainable" in summary[layer]: - if summary[layer]["trainable"] == True: - trainable_params += summary[layer]["nb_params"] summary_str += line_new + "\n" + params = list(model.parameters()) + total_params = sum(p.numel() for p in params) + trainable_params = sum(p.numel() for p in params if p.requires_grad) + frozen_params = total_params - trainable_params + # assume 4 bytes/number (float on cuda). total_input_size = abs( np.prod(sum(input_size, ())) * batch_size * 4.0 / (1024**2.0) @@ -169,9 +184,7 @@ def hook(module, input, output): summary_str += "=" * (max_layer_len + 43) + "\n" summary_str += "Total params: {0:,}".format(total_params) + "\n" summary_str += "Trainable params: {0:,}".format(trainable_params) + "\n" - summary_str += ( - "Non-trainable params: {0:,}".format(total_params - trainable_params) + "\n" - ) + summary_str += "Non-trainable params: {0:,}".format(frozen_params) + "\n" summary_str += "-" * (max_layer_len + 43) + "\n" summary_str += "Input size (MB): %0.2f" % total_input_size + "\n" summary_str += "Forward/backward pass size (MB): %0.2f" % total_output_size + "\n"