Fix adversarial image visualizer with canonical batches (#227)

* Fetch canonical input and target.

* Invoke Adversary.forward().

* Fix test.

mart/callbacks/visualizer.py

@@ -27,15 +27,14 @@ def __init__(self, folder):
             os.makedirs(self.folder)
 
     def on_train_batch_end(self, trainer, model, outputs, batch, batch_idx):
-        # Save input and target for on_train_end
-        self.input = batch["input"]
-        self.target = batch["target"]
+        # Save canonical input and target for on_train_end
+        self.input = batch[0]
+        self.target = batch[1]
 
     def on_train_end(self, trainer, model):
         # FIXME: We should really just save this to outputs instead of recomputing adv_input
         with torch.no_grad():
-            perturbation = model.perturber(input=self.input, target=self.target)
-            adv_input = model.composer(perturbation, input=self.input, target=self.target)
+            adv_input, _target = model(self.input, self.target)
 
         for img, tgt in zip(adv_input, self.target):
             fname = tgt["file_name"]

tests/test_visualizer.py

@@ -19,19 +19,21 @@ def test_visualizer_run_end(input_data, target_data, perturbation, tmp_path):
     target_list = [target_data]
 
     # simulate an addition perturbation
-    def perturb(input):
+    def perturb(input, target):
         result = [sample + perturbation for sample in input]
-        return result
+        return result, target
 
+    adversary = Mock(spec=Adversary, side_effect=perturb)
     trainer = Mock()
-    model = Mock(composer=Mock(return_value=perturb(input_list)))
     outputs = Mock()
-    batch = {"input": input_list, "target": target_list}
-    adversary = Mock(spec=Adversary, side_effect=perturb)
+    target_model = Mock()
+
+    # Canonical batch in Adversary.
+    batch = (input_list, target_list, target_model)
 
     visualizer = PerturbedImageVisualizer(folder)
-    visualizer.on_train_batch_end(trainer, model, outputs, batch, 0)
-    visualizer.on_train_end(trainer, model)
+    visualizer.on_train_batch_end(trainer, adversary, outputs, batch, 0)
+    visualizer.on_train_end(trainer, adversary)
 
     # verify that the visualizer created the JPG file
     expected_output_path = folder / target_data["file_name"]