IMPROVEMENT: handle models with grayscale input (#31)

* Handle models with grayscale input

* Fix spacing

* Fix linting

* Bump version

flashtorch/__init__.py

@@ -1 +1 @@
-__version__ = '0.1.2'
+__version__ = '0.1.3'

flashtorch/saliency/backprop.py

@@ -218,8 +218,7 @@ def _record_gradients(module, grad_in, grad_out):
                 self.gradients = grad_in[0]
 
         for _, module in self.model.named_modules():
-            if isinstance(module, nn.modules.conv.Conv2d) and \
-                    module.in_channels == 3:
+            if isinstance(module, nn.modules.conv.Conv2d):
                 module.register_backward_hook(_record_gradients)
                 break
 

setup.py

@@ -26,7 +26,7 @@
 DOCLINES = (__doc__ or '').split("\n")
 long_description = "\n".join(DOCLINES[2:])
 
-version = '0.1.2'
+version = '0.1.3'
 
 setup(
     name='flashtorch',

tests/test_backprop.py

@@ -3,6 +3,7 @@
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 import torchvision.models as models
 
@@ -59,6 +60,21 @@ def make_expected_gradient_target(top_class):
     return target
 
 
+class CnnGrayscale(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=3, padding=1)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(64, 10, kernel_size=3, stride=3, padding=1)
+        self.fc1 = nn.Linear(10 * 25 * 25, 10)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+
+        return F.softmax(self.fc1(x.view(-1, 10 * 25 * 25)), dim=1)
+
+
 #################
 # Test fixtures #
 #################
@@ -69,6 +85,11 @@ def model():
     return models.alexnet()
 
 
+@pytest.fixture
+def model_grayscale():
+    return CnnGrayscale()
+
+
 ##############
 # Test cases #
 ##############
@@ -162,6 +183,16 @@ def test_calc_gradients_of_top_class_if_prediction_is_wrong(mocker, model):
     assert torch.all(kwargs['gradient'].eq(target))
 
 
+def test_handle_greyscale_input(mocker, model_grayscale):
+    backprop = Backprop(model_grayscale)
+
+    input_ = torch.zeros([1, 1, 224, 224], requires_grad=True)
+
+    gradients = backprop.calculate_gradients(input_)
+
+    assert gradients.shape == (1, 224, 224)
+
+
 def test_return_max_across_color_channels_if_specified(mocker, model):
     backprop = Backprop(model)