Merge pull request #106 from sksq96/issue_9

can now handle single layers, for #9

Author: Naireen

.gitignore

@@ -1 +1,2 @@
 __pycache__
+*.pyc

torchsummary/tests/unit_tests/torchsummary_test.py

@@ -1,7 +1,7 @@
 import unittest
 from torchsummary import summary
 from torchsummary.tests.test_models.test_model import SingleInputNet, MultipleInputNet
-
+import torch
 
 class torchsummaryTests(unittest.TestCase):
     def test_single_input(self):
@@ -19,5 +19,20 @@ def test_multiple_input(self):
         self.assertEqual(total_params, 31120)
         self.assertEqual(trainable_params, 31120)
 
+    def test_single_layer_network(self):
+        model = torch.nn.Linear(2, 5)
+        input = (1, 2)
+        total_params, trainable_params = summary(model, input, device="cpu")
+        self.assertEqual(total_params, 15)
+        self.assertEqual(trainable_params, 15)
+
+    def test_single_layer_network_on_gpu(self):
+        model = torch.nn.Linear(2, 5)
+        model.cuda()
+        input = (1, 2)
+        total_params, trainable_params = summary(model, input, device="cuda")
+        self.assertEqual(total_params, 15)
+        self.assertEqual(trainable_params, 15)
+
 if __name__ == '__main__':
     unittest.main(buffer=True)

torchsummary/torchsummary.py

@@ -37,7 +37,6 @@ def hook(module, input, output):
         if (
             not isinstance(module, nn.Sequential)
             and not isinstance(module, nn.ModuleList)
-            and not (module == model)
         ):
             hooks.append(module.register_forward_hook(hook))
 
@@ -90,6 +89,7 @@ def hook(module, input, output):
             "{0:,}".format(summary[layer]["nb_params"]),
         )
         total_params += summary[layer]["nb_params"]
+
         total_output += np.prod(summary[layer]["output_shape"])
         if "trainable" in summary[layer]:
             if summary[layer]["trainable"] == True:
@@ -99,7 +99,7 @@ def hook(module, input, output):
     # 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.numpy() * 4. / (1024 ** 2.))
+    total_params_size = abs(total_params * 4. / (1024 ** 2.))
     total_size = total_params_size + total_output_size + total_input_size
 
     print("================================================================")