Repository: mlcommons/GaNDLF

Issue:
Add attention maps
**Is your feature request related to a problem? Please describe.**
Interpretability in DL training is critical and would be very nice to have in GaNDLF.

**Describe the solution you'd like**
Integrate [M3d-Cam](https://github.com/MECLabTUDA/M3d-Cam) into the training/inference process.

**Describe alternatives you've considered**
Open to other suggestions.

**Additional context**
Also requested by @gastouna



Discussion:
Hi,

I successfully integraded medcam into the validation section of the training already.
Currently there are probably some issues with assembling the attention patches with the `torchio.inference.GridAggregator`, but I am not entirely sure. These issues could also simply originate from the fact that the model is untrained.
Do you have a fully trained 2D segmentation dataset, model and config yaml at hand that you could send me?

Best
Karol
> I successfully integraded medcam into the validation section of the training already.

Awesome news!!!

> These issues could also simply originate from the fact that the model is untrained.

Could you elaborate? I think the `torchio.inference.GridAggregator` does not care if a model is completely trained or not, as long as it is fed patch locations [[ref](https://github.com/CBICA/GaNDLF/blob/refactor/GANDLF/training_loop.py#L276)].

> Do you have a fully trained 2D segmentation dataset, model and config yaml at hand that you could send me?

I do not have any such example on hand, perhaps @Geeks-Sid does? 
Here is an example of the attention superimposed over the original image input:

![attention_map](https://user-images.githubusercontent.com/3471895/113307605-9cd11600-9305-11eb-830a-885df2f31263.png)

The model is of course not able to learn from just a few images what it should pay attention to, so it is more or less random. However, you can clearly see borders and different normalizations. My hope is that this will go away when the model is trained. I integrated medcam into the nnUNet for example as well which is also patch-based and had no such problems (but I only tested on the finished models). But it could also be the case that there is some problem with normalization or interpolation.
Ah okay. That actually looks great! I'll ping @Geeks-Sid if he has a fully trained model.

Resolving patch:
diff --git a/GANDLF/training_loop.py b/GANDLF/training_loop.py
--- a/GANDLF/training_loop.py
+++ b/GANDLF/training_loop.py
@@ -17,6 +17,7 @@
 from GANDLF.parameterParsing import get_model, get_optimizer, get_scheduler
 from GANDLF.utils import get_date_time, send_model_to_device
 from GANDLF.data.ImagesFromDataFrame import ImagesFromDataFrame
+from medcam import medcam
 
 os.environ["TORCHIO_HIDE_CITATION_PROMPT"] = "1"  # hides torchio citation request
 
@@ -89,14 +90,23 @@ def step(model, image, label, params):
             output = model(image)
     else:
         output = model(image)
+
+    if "medcam_enabled" in params and params["medcam_enabled"]:
+        output, attention_map = output
+
     # print("Output shape : ", output.shape, flush=True)
     # one-hot encoding of 'output' will probably be needed for segmentation
     loss, metric_output = get_loss_and_metrics(label, output, params)
     
     if params["model"]["dimension"] == 2:
         output = torch.unsqueeze(output, -1)
+        if "medcam_enabled" in params and params["medcam_enabled"]:
+            attention_map = torch.unsqueeze(attention_map, -1)
 
-    return loss, metric_output, output
+    if not ("medcam_enabled" in params and params["medcam_enabled"]):
+        return loss, metric_output, output
+    else:
+        return loss, metric_output, output, attention_map
 
 
 def train_network(model, train_dataloader, optimizer, params):
@@ -227,6 +237,9 @@ def validate_network(model, valid_dataloader, scheduler, params):
 
     # Set the model to valid
     model.eval()
+    if "medcam_enabled" in params:
+        model.enable_medcam()
+        params["medcam_enabled"] = True
     for batch_idx, (subject) in enumerate(valid_dataloader):
         
         # constructing a new dict because torchio.GridSampler requires torchio.Subject, which requires torchio.Image to be present in initial dict, which the loader does not provide
@@ -252,6 +265,8 @@ def validate_network(model, valid_dataloader, scheduler, params):
         grid_sampler = torchio.inference.GridSampler(torchio.Subject(subject_dict), params['patch_size'])
         patch_loader = torch.utils.data.DataLoader(grid_sampler, batch_size=1)
         aggregator = torchio.inference.GridAggregator(grid_sampler)
+        if "medcam_enabled" in params:
+            attention_map_aggregator = torchio.inference.GridAggregator(grid_sampler)
         
         is_segmentation = True
         output_prediction = 0 # this is used for regression/classification
@@ -271,7 +286,11 @@ def validate_network(model, valid_dataloader, scheduler, params):
             label = label.to(params["device"])
             # print("Validation :", label.shape, image.shape, flush=True)
             patch_location = patches_batch[torchio.LOCATION]
-            loss, calculated_metrics, output = step(model, image, label, params)
+            result = step(model, image, label, params)
+            if "medcam_enabled" not in params:
+                loss, calculated_metrics, output = result
+            else:
+                loss, calculated_metrics, output, attention_map = result
             if is_segmentation:
                 aggregator.add_batch(output, patch_location)
             else:
@@ -279,6 +298,8 @@ def validate_network(model, valid_dataloader, scheduler, params):
                     output_prediction += output.cpu().data.item()# this probably needs customization for classification (majority voting or median, perhaps?)
                 else:
                     output_prediction += output
+            if "medcam_enabled" in params:
+                attention_map_aggregator.add_batch(attention_map, patch_location)
         
         if is_segmentation:
             output_prediction = aggregator.get_output_tensor()
@@ -286,6 +307,11 @@ def validate_network(model, valid_dataloader, scheduler, params):
         else:
             output_prediction = output_prediction / len(patch_loader) # final regression output
 
+        if "medcam_enabled" in params:
+            attention_map = attention_map_aggregator.get_output_tensor()
+            for i in range(len(attention_map)):
+                model.save_attention_map(attention_map[i].squeeze(), raw_input=image[i].squeeze(-1))
+
         ## this is currently broken
         # final_loss, final_metric = get_loss_and_metrics(label_ground_truth, output_prediction, params)
         # print("Full image validation:: Loss: ", final_loss, "; Metric: ", final_metric, flush=True)
@@ -307,6 +333,10 @@ def validate_network(model, valid_dataloader, scheduler, params):
                         total_epoch_valid_metric[metric] / len(valid_dataloader),
                     )
 
+    if "medcam_enabled" in params:
+        model.disable_medcam()
+        params["medcam_enabled"] = False
+
     average_epoch_valid_loss = total_epoch_valid_loss / len(valid_dataloader)
     for metric in params["metrics"]:
         average_epoch_valid_metric[metric] = total_epoch_valid_metric[metric] / len(
@@ -489,6 +519,10 @@ def training_loop(
         model, amp=params["model"]["amp"], device=params["device"], optimizer=optimizer
     )
 
+    if "medcam" in params:
+        model = medcam.inject(model, output_dir=os.path.join(output_dir, "attention_maps", params["medcam"]["backend"]), backend=params["medcam"]["backend"], layer=params["medcam"]["layer"], save_maps=False, return_attention=True, enabled=False)
+        params["medcam_enabled"] = False
+
     # Setup a few variables for tracking
     best_loss = 1e7
     patience = 0
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -30,7 +30,8 @@
   'pytest',
   'pytest-azurepipelines',
   'coverage',
-  'psutil'
+  'psutil',
+  'medcam'
 ]
 
 setup(


Test patch: