feat: device-agnostic confidence

Author: jannisborn

paccmann_predictor/utils/interpret.py

@@ -1,20 +1,30 @@
+from typing import Tuple
+
 import torch
-from torch import nn
+from torch import Tensor, nn
+
+from .utils import get_device
 
 # We use standard deviation to measure uncertainity since entropy is not
 # defined for continuous variables and differential entropy is not ideal.
 # In case all predictions are identical, std is 0. If 50% are 0 and 50% are
 # one, it is maximal, i.e. 0.5.
 MAX_STD = 0.5
-MIN_STD = 0.
+MIN_STD = 0.0
+
+DEVICE = get_device()
+
+
+def map_to_device(inputs: Tuple[Tensor, ...]) -> Tuple[Tensor, ...]:
+    return tuple(x.to(DEVICE) for x in inputs)
 
 
 def monte_carlo_dropout(
-    model, regime='loader', loader=None, tensors=None, repetitions=20
+    model, regime="loader", loader=None, tensors=None, repetitions=20
 ):
     """
     Attempts to approximate epistemic uncertainity through MC dropout.
-    Performs Monte Carlo dropout for a given model and returns a list of 
+    Performs Monte Carlo dropout for a given model and returns a list of
     sample-wise confidence estimates.
     This method can be used in two regimes, either by passing a dataloader
     or by passing a tensor with the raw input to the model.
@@ -25,7 +35,7 @@ def monte_carlo_dropout(
 
 
     Arguments:
-        model (torch.nn.Module): The torch network to be investigated. 
+        model (torch.nn.Module): The torch network to be investigated.
             NOTE: Model is assumed to return either a single tensor of
             predictions or a n-tupel with the first part being a tensor
             of predictions. They need to be [0, 1] where 0 and 1 represent
@@ -50,7 +60,7 @@ def monte_carlo_dropout(
             Contains the averaged predictions across all MC  dropout estimates.
     """
 
-    if regime != 'loader' and regime != 'tensors':
+    if regime != "loader" and regime != "tensors":
         raise ValueError("Choose regime from {'loader', 'tensors'}")
 
     # Activate dropout layers while keeping other rest in eval mode.
@@ -61,53 +71,48 @@ def enable_dropout(m):
     model.eval()
     model.apply(enable_dropout)
 
-    if regime == 'loader':
+    if regime == "loader":
 
         # Error handling
-        if not isinstance(
-            loader.sampler, torch.utils.data.sampler.SequentialSampler
-        ):
+        if not isinstance(loader.sampler, torch.utils.data.sampler.SequentialSampler):
             raise AttributeError(
-                'Data loader does not use sequential sampling. Consider set'
-                'ting shuffle=False when instantiating the data loader.'
+                "Data loader does not use sequential sampling. Consider set"
+                "ting shuffle=False when instantiating the data loader."
             )
 
         # Run over all batches in the loader
 
         def call_fn():
             preds = []
-            for ind, inputs in enumerate(loader):
+            for inputs in loader:
                 # inputs is a tuple with the last element being the labels
                 # outs can be a n-tuple returned by the model
-                outs = model(*inputs[:-1])
-                preds.append(outs[0] if isinstance(outs, tuple) else outs)
+                outs = model(*map_to_device(inputs[:-1]))
+                preds.append(
+                    outs[0].detach().cpu()
+                    if isinstance(outs, tuple)
+                    else outs.detach().cpu()
+                )
 
             return torch.cat(preds)
 
-    elif regime == 'tensors':
+    elif regime == "tensors":
 
-        if (
-            not isinstance(tensors, tuple)
-            and not isinstance(tensors, torch.Tensor)
-        ):
-            raise ValueError('Tensor needs to either tuple or torch.Tensor')
+        if not isinstance(tensors, tuple) and not isinstance(tensors, torch.Tensor):
+            raise ValueError("Tensor needs to either tuple or torch.Tensor")
 
-        inputs = tensors if isinstance(tensors, tuple) else (tensors, )
+        inputs = tensors if isinstance(tensors, tuple) else (tensors,)
 
         def call_fn():
-            outs = model(*inputs)
+            outs = model(*map_to_device(inputs))
             return outs[0] if isinstance(outs, tuple) else outs
 
     with torch.no_grad():
-        predictions = [
-            torch.unsqueeze(call_fn(), -1) for _ in range(repetitions)
-        ]
+        predictions = [torch.unsqueeze(call_fn(), -1) for _ in range(repetitions)]
     predictions = torch.cat(predictions, dim=-1)
 
     # Scale confidences to [0, 1]
-    confidences = -1 * (
-        (predictions.std(dim=-1) - MIN_STD) / (MAX_STD - MIN_STD)
-    ) + 1
+    confidences = -1 * ((predictions.std(dim=-1) - MIN_STD) / (MAX_STD - MIN_STD)) + 1
 
     model.eval()
 
@@ -116,12 +121,12 @@ def call_fn():
 
 def test_time_augmentation(
     model,
-    regime='loader',
+    regime="loader",
     loader=None,
     tensors=None,
     repetitions=20,
     augmenter=None,
-    tensors_to_augment=None
+    tensors_to_augment=None,
 ):
     """
     Attempts to measure aleatoric uncertainity through augmentation during test
@@ -135,7 +140,7 @@ def test_time_augmentation(
     classification like MNIST.
 
     Arguments:
-        model (torch.nn.Module): The torch network to be investigated. 
+        model (torch.nn.Module): The torch network to be investigated.
             NOTE: Model is assumed to return either a single tensor of
             predictions or a n-tupel with the first part being a tensor
             of predictions. They need to be [0, 1] where 0 and 1 represent
@@ -173,86 +178,80 @@ def test_time_augmentation(
             Contains the averaged predictions across estimates.
     """
 
-    if regime != 'loader' and regime != 'tensors':
+    if regime != "loader" and regime != "tensors":
         raise ValueError("Choose regime from {'loader', 'tensors'}")
 
     model.eval()
 
-    if regime == 'loader':
+    if regime == "loader":
 
         # Error handling
-        if not isinstance(
-            loader.sampler, torch.utils.data.sampler.SequentialSampler
-        ):
+        if not isinstance(loader.sampler, torch.utils.data.sampler.SequentialSampler):
             raise AttributeError(
-                'Data loader does not use sequential sampling. Consider set'
-                'ting shuffle=False when instantiating the data loader.'
+                "Data loader does not use sequential sampling. Consider set"
+                "ting shuffle=False when instantiating the data loader."
             )
 
         # Run over all batches in the loader
 
         def call_fn():
             preds = []
-            for ind, inputs in enumerate(loader):
+            for inputs in loader:
                 # inputs is a tuple with the last element being the labels
                 # outs can be a n-tuple returned by the model
-                outs = model(*inputs[:-1])
+                outs = model(*map_to_device(inputs[:-1]))
                 preds.append(outs[0] if isinstance(outs, tuple) else outs)
 
             return torch.cat(preds)
 
-    elif regime == 'tensors':
+    elif regime == "tensors":
 
-        if (
-            not isinstance(tensors, tuple)
-            and not isinstance(tensors, torch.Tensor)
-        ):
-            raise ValueError('Tensor needs to either tuple or torch.Tensor')
-        if (
-            not isinstance(tensors_to_augment, list)
-            and not isinstance(tensors_to_augment, int)
+        if not isinstance(tensors, tuple) and not isinstance(tensors, torch.Tensor):
+            raise ValueError("Tensor needs to either tuple or torch.Tensor")
+        if not isinstance(tensors_to_augment, list) and not isinstance(
+            tensors_to_augment, int
         ):
-            raise ValueError('tensors_to_augment needs to be list or int')
+            raise ValueError("tensors_to_augment needs to be list or int")
 
         # Convert input to common formats (tuples and lists)
         tensors_to_augment = (
             [tensors_to_augment]
-            if isinstance(tensors_to_augment, int) else tensors_to_augment
+            if isinstance(tensors_to_augment, int)
+            else tensors_to_augment
         )
-        inputs = tensors if isinstance(tensors, tuple) else (tensors, )
-        aug_fns = augmenter if isinstance(augmenter, tuple) else (augmenter, )
+        inputs = tensors if isinstance(tensors, tuple) else (tensors,)
+        aug_fns = augmenter if isinstance(augmenter, tuple) else (augmenter,)
 
         # Error handling
         if not len(aug_fns) == len(tensors_to_augment):
             raise ValueError(
-                'Provide one augmenter for each tensor you want to augment.'
+                "Provide one augmenter for each tensor you want to augment."
             )
         if max(tensors_to_augment) > len(inputs):
             raise ValueError(
-                'tensors_to_augment should be indexes to the tensors used for '
-                f'augmentation. {max(tensors_to_augment)} is larger than '
-                f'length of inputs ({len(inputs)}).'
+                "tensors_to_augment should be indexes to the tensors used for "
+                f"augmentation. {max(tensors_to_augment)} is larger than "
+                f"length of inputs ({len(inputs)})."
             )
 
         def call_fn():
             # Perform augmentation on all designated functions
             augmented_inputs = [
-                aug_fns[tensors_to_augment[tensors_to_augment == ind]](tensor)
-                if ind in tensors_to_augment else tensor
+                (
+                    aug_fns[tensors_to_augment[tensors_to_augment == ind]](tensor)
+                    if ind in tensors_to_augment
+                    else tensor
+                )
                 for ind, tensor in enumerate(tensors)
             ]
-            outs = model(*augmented_inputs)
+            outs = model(*map_to_device(augmented_inputs))
             return outs[0] if isinstance(outs, tuple) else outs
 
     with torch.no_grad():
-        predictions = [
-            torch.unsqueeze(call_fn(), -1) for _ in range(repetitions)
-        ]
+        predictions = [torch.unsqueeze(call_fn(), -1) for _ in range(repetitions)]
     predictions = torch.cat(predictions, dim=-1)
 
     # Scale confidences to [0, 1]
-    confidences = -1 * (
-        (predictions.std(dim=-1) - MIN_STD) / (MAX_STD - MIN_STD)
-    ) + 1
+    confidences = -1 * ((predictions.std(dim=-1) - MIN_STD) / (MAX_STD - MIN_STD)) + 1
 
     return torch.clamp(confidences, min=0), torch.mean(predictions, -1)