Backprop.visualize passes use_gpu flag (#25)

* Backprop.visualize passes use_gpu flag * Bump version
MisaOgura · Jan 2, 2020 · d48b391 · d48b391
1 parent aa57862
commit d48b391
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 4 deletions.
diff --git a/flashtorch/__init__.py b/flashtorch/__init__.py
@@ -1 +1 @@
-__version__ = '0.1.1'
+__version__ = '0.1.2'
diff --git a/flashtorch/saliency/backprop.py b/flashtorch/saliency/backprop.py
@@ -165,11 +165,13 @@ def visualize(self, input_, target_class, guided=False, use_gpu=False,
 
         gradients = self.calculate_gradients(input_,
                                              target_class,
-                                             guided=guided)
+                                             guided=guided,
+                                             use_gpu=use_gpu)
         max_gradients = self.calculate_gradients(input_,
                                                  target_class,
                                                  guided=guided,
-                                                 take_max=True)
+                                                 take_max=True,
+                                                 use_gpu=use_gpu)
 
         # Setup subplots
 

diff --git a/setup.py b/setup.py
@@ -26,7 +26,7 @@
 DOCLINES = (__doc__ or '').split("\n")
 long_description = "\n".join(DOCLINES[2:])
 
-version = '0.1.1'
+version = '0.1.2'
 
 setup(
     name='flashtorch',

diff --git a/tests/test_backprop.py b/tests/test_backprop.py
@@ -213,6 +213,45 @@ def test_warn_when_prediction_is_wrong(mocker, model):
         backprop.calculate_gradients(input_, target_class)
 
 
+# Test visualize method
+
+
+def test_visualize_calls_calculate_gradients_twice(mocker, model):
+    backprop = Backprop(model)
+    mocker.spy(backprop, 'calculate_gradients')
+
+    top_class = 5
+    target_class = 5
+    input_ = torch.zeros([1, 3, 224, 224])
+
+    target = make_expected_gradient_target(top_class)
+
+    mock_output = make_mock_output(mocker, model, target_class)
+
+    backprop.visualize(input_, target_class, use_gpu=True)
+
+    assert backprop.calculate_gradients.call_count == 2
+
+
+def test_visualize_passes_gpu_flag(mocker, model):
+    backprop = Backprop(model)
+    mocker.spy(backprop, 'calculate_gradients')
+
+    top_class = 5
+    target_class = 5
+    input_ = torch.zeros([1, 3, 224, 224])
+
+    target = make_expected_gradient_target(top_class)
+
+    mock_output = make_mock_output(mocker, model, target_class)
+
+    backprop.visualize(input_, target_class, use_gpu=True)
+
+    _, _, kwargs = backprop.calculate_gradients.mock_calls[0]
+
+    assert kwargs['use_gpu']
+
+
 # Test compatibilities with torchvision models