pass in device and dtypes explicitly for #102 among others

Author: Naireen

torchsummary/tests/test_models/test_model.py

@@ -35,3 +35,22 @@ def forward(self, x1, x2):
         x2 = self.fc2b(x2)
         x = torch.cat((x1, x2), 0)
         return F.log_softmax(x, dim=1)
+
+class MultipleInputNetDifferentDtypes(nn.Module):
+    def __init__(self):
+        super(MultipleInputNetDifferentDtypes, self).__init__()
+        self.fc1a = nn.Linear(300, 50)
+        self.fc1b = nn.Linear(50, 10)
+
+        self.fc2a = nn.Linear(300, 50)
+        self.fc2b = nn.Linear(50, 10)
+
+    def forward(self, x1, x2):
+        x1 = F.relu(self.fc1a(x1))
+        x1 = self.fc1b(x1)
+        x2 = x2.type(torch.FloatTensor)
+        x2 = F.relu(self.fc2a(x2))
+        x2 = self.fc2b(x2)
+        # set x2 to FloatTensor
+        x = torch.cat((x1, x2), 0)
+        return F.log_softmax(x, dim=1)

torchsummary/tests/unit_tests/torchsummary_test.py

@@ -1,6 +1,6 @@
 import unittest
 from torchsummary import summary
-from torchsummary.tests.test_models.test_model import SingleInputNet, MultipleInputNet
+from torchsummary.tests.test_models.test_model import SingleInputNet, MultipleInputNet, MultipleInputNetDifferentDtypes
 import torch
 
 class torchsummaryTests(unittest.TestCase):
@@ -30,9 +30,18 @@ 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")
+        total_params, trainable_params = summary(model, input, device="cuda:0")
         self.assertEqual(total_params, 15)
         self.assertEqual(trainable_params, 15)
 
+    def test_multiple_input_types(self):
+        model = MultipleInputNetDifferentDtypes()
+        input1 = (1, 300)
+        input2 = (1, 300)
+        dtypes = [torch.FloatTensor, torch.LongTensor]
+        total_params, trainable_params = summary(model, [input1, input2], device="cpu", dtypes=dtypes)
+        self.assertEqual(total_params, 31120)
+        self.assertEqual(trainable_params, 31120)
+
 if __name__ == '__main__':
     unittest.main(buffer=True)

torchsummary/torchsummary.py

@@ -6,7 +6,9 @@
 import numpy as np
 
 
-def summary(model, input_size, batch_size=-1, device="cuda"):
+def summary(model, input_size, batch_size=-1, device=torch.device('cuda:0'), dtypes=None):
+    if dtypes == None:
+        dtypes = [torch.FloatTensor]*len(input_size)
 
     def register_hook(module):
 
@@ -40,24 +42,13 @@ def hook(module, input, output):
         ):
             hooks.append(module.register_forward_hook(hook))
 
-    device = device.lower()
-    assert device in [
-        "cuda",
-        "cpu",
-    ], "Input device is not valid, please specify 'cuda' or 'cpu'"
-
-    if device == "cuda" and torch.cuda.is_available():
-        dtype = torch.cuda.FloatTensor
-    else:
-        dtype = torch.FloatTensor
-
     # multiple inputs to the network
     if isinstance(input_size, tuple):
         input_size = [input_size]
 
+
     # batch_size of 2 for batchnorm
-    x = [torch.rand(2, *in_size).type(dtype) for in_size in input_size]
-    # print(type(x[0]))
+    x = [ torch.rand(2, *in_size).type(dtype).to(device=device) for in_size, dtype in zip(input_size, dtypes)]
 
     # create properties
     summary = OrderedDict()