From ad3f949faa45a7e107851e38b0c3597dc8669d93 Mon Sep 17 00:00:00 2001 From: Soumyadeep Bhattacharjee Date: Sat, 19 Mar 2022 03:11:23 +0100 Subject: [PATCH] Added OrderedDict as an output datatype Some pre-trained models like the deeplabv3_resnet50 returns an OrderedDict containing the semantic mask and the auxillary loss as two separate Tensors under the keys 'out' and 'aux' (https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101/). Usage of torchsummary with such models was not compatible since OrderedDict had not been considered as a potential output type. This issue has now been fixed. --- torchsummary/torchsummary.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/torchsummary/torchsummary.py b/torchsummary/torchsummary.py index 1ed065f..6b92925 100644 --- a/torchsummary/torchsummary.py +++ b/torchsummary/torchsummary.py @@ -33,6 +33,14 @@ def hook(module, input, output): summary[m_key]["output_shape"] = [ [-1] + list(o.size())[1:] for o in output ] + elif isinstance(output,OrderedDict): + keys = list(output.keys()) + if 'out' in keys: + key='out' + else: + key = keys[-1] + summary[m_key]["output_shape"] = list(output[key].size()) + summary[m_key]["output_shape"][0] = batch_size else: summary[m_key]["output_shape"] = list(output.size()) summary[m_key]["output_shape"][0] = batch_size