Aleatoric and Epistemic Uncertainty Engine

[danieltudosiu]

Is your feature request related to a problem? Please describe.
Create an inferer for aleatoric and epistemic uncertainty.

Describe the solution you'd like
The inferer can be based on the TestTimeAugmentation that is already implemented and MC-Dropout uncertainty.

[danieltudosiu]

Closing as it is not possible as Inferer it needs to be an Engine

[danieltudosiu]

Upon further inspection of the TestTimeAugmentation could be transformed into a peculiar Engine.

The engine would differ from the usual ones by not taking in a DataLoader but a Dict of unique subjects that internally they will be passed to a Dataset constructor similar to TestTimeAugmentation logic here. Furthermore, the Engine will take all the other required arguments for the creation of the Dataset and Dataloader.

Lastly, the Engine will be an Evaluator style engine but it will keep the model in training mode if MC-Dropout uncertainty is required or put it to eval otherwise.

My only worry is that the Workflow requires the passing of a dataloader which invalidates the proposed approach. If we are to decouple that requirement we could create an UncertaintyEngine for MONAI.

Alternatively, we could try and make it an Inferer, but that would cause issues with the usage of a SlidingWindowInferer and UncertaintyInferer and I think their compatibility is required for low VRAM medical imagining use cases.

I will internally work on this but can't promise anything in terms of PR. More than happy to collaborate on it.

[danieltudosiu]

If somebody is picking this up, here is my go at it.

class TestTimeAugmentation:
    """
    Class for performing test time augmentations. This will pass the same image through the network multiple times.

    The user passes transform(s) to be applied to each realisation, and provided that at least one of those transforms
    is random, the network's output will vary. Provided that inverse transformations exist for all supplied spatial
    transforms, the inverse can be applied to each realisation of the network's output. Once in the same spatial
    reference, the results can then be combined and metrics computed.

    Test time augmentations are a useful feature for computing network uncertainty, as well as observing the network's
    dependency on the applied random transforms.

    To calculate the VVC you can use the following code:
        vvc = (np.std(output) / np.mean(output)).item()

    Reference:
        Wang et al.,
        Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional
        neural networks,
        https://doi.org/10.1016/j.neucom.2019.01.103

    Args:
        transform: transform (or composed) to be applied to each realisation. At least one transform must be of type
            `Randomizable`. All random transforms must be of type `InvertibleTransform`.
        network: network to evaluate in the evaluator, should be regular PyTorch `torch.nn.Module`.
        aleatoric_realisations: number of realisations to use for computing aleatoric uncertainty.
        epistemic_realisations: number of realisations to use for computing epistemic uncertainty.
        num_workers: how many subprocesses to use for data.
        inferer: inference method that execute model forward on input data, like: SlidingWindow, etc.
        device: device on which to perform inference.
        image_key: key used to extract image from input dictionary.
        orig_key: the key of the original input data in the dict. will get the applied transform information
            for this input data, then invert them for the expected data with `image_key`.
        orig_meta_keys: the key of the meta data of original input data, will get the `affine`, `data_shape`, etc.
            the meta data is a dictionary object which contains: filename, original_shape, etc.
            if None, will try to construct meta_keys by `{orig_key}_{meta_key_postfix}`.
        meta_key_postfix: use `key_{postfix}` to to fetch the meta data according to the key data,
            default is `meta_dict`, the meta data is a dictionary object.
            For example, to handle key `image`,  read/write affine matrices from the
            metadata `image_meta_dict` dictionary's `affine` field.
            this arg only works when `meta_keys=None`.
        return_full_data: normally, metrics are returned (mode, mean, std, vvc). Setting this flag to `True` will return
            the full data. Dimensions will be same size as when passing a single image through `inferrer_fn`, with a
            dimension appended equal in size to `num_examples` (N), i.e., `[N,C,H,W,[D]]`.
        progress: whether to display a progress bar.
    """

    def __init__(
        self,
        transform: InvertibleTransform,
        network: Callable[..., torch.Tensor],
        post_inferer_transform: Optional[Transform],
        output_transform: Optional[Transform],
        aleatoric_realisations: int = 5,
        epistemic_realisations: int = 5,
        num_workers: int = 1,
        inferer: Optional[Inferer] = None,
        device: Union[str, torch.device] = "gpu:0",
        image_key=DataSetKeys.RGB.value,
        orig_key=CommonKeys.LABEL,
        nearest_interp: bool = True,
        orig_meta_keys: Optional[str] = None,
        meta_key_postfix="meta_dict",
        return_full_data: bool = False,
        progress: bool = False,
    ) -> None:
        self.transform = transform
        self.network = network
        self.aleatoric_realisations = aleatoric_realisations
        self.epistemic_realisations = epistemic_realisations
        self.num_workers = num_workers
        self.inferer = SimpleInferer() if inferer is None else inferer
        self.device = device
        self.image_key = image_key
        self.orig_key = orig_key
        self.nearest_interp = nearest_interp
        self.orig_meta_keys = orig_meta_keys
        self.meta_key_postfix = meta_key_postfix
        self.return_full_data = return_full_data
        self.progress = progress
        self._decolate = Decollated(keys=None, detach=True)
        # check that the transform has at least one random component, and that all random transforms are invertible
        self._check_transforms()

        self._post_inferer_transform = post_inferer_transform
        self._output_transform = output_transform

    def _check_transforms(self):
        """
            Check if self.transform is at least 1 random transform, and all random transforms should be invertible.
        """

        # Make sure the transformations are a list
        ts = (
            [self.transform]
            if not isinstance(self.transform, Compose)
            else self.transform.transforms
        )

        # Create a binary mask of random and invertible transforms
        randoms = np.array([isinstance(t, Randomizable) for t in ts])
        invertibles = np.array([isinstance(t, InvertibleTransform) for t in ts])

        # Check if at least one transformation is Random
        if sum(randoms) == 0:
            raise RuntimeError(
                "Requires a `Randomizable` transform or a `Compose` containing at least one `Randomizable` transform."
            )

        # Check if all random transformations are invertible
        for r, i in zip(randoms, invertibles):
            if r and not i:
                warnings.warn(
                    f"Not all applied random transform(s) are invertible. Problematic transform: {type(r).__name__}"
                )

    @torch.no_grad()
    def __call__(
        self, data: List[Dict[str, str]]
    ) -> Union[Tuple[np.ndarray, np.ndarray, np.ndarray, float], np.ndarray]:
        """
        Args:
            data: dictionary data to be processed.

        Returns:
            - if `return_full_data==False`: mode, mean, std, vvc. The mode, mean and standard deviation are calculated
                across 'num_examples` outputs at each voxel. The volume variation coefficient (VVC) is `std/mean` across
                the whole output, including `num_examples`. See original paper for clarification.
            - if `return_full_data==True`: data is returned as-is after applying the `inferrer_fn` and then
                concatenating across the first dimension containing `num_examples`. This allows the user to perform
                their own analysis if desired.
        """
        # Set the network on training if we are estimating the epistemic uncertainty
        if self.epistemic_realisations > 1:
            self.network.train()

        # Apply uncertainty logic for each subject
        for subject in data:
            # Copy the data num_examples times
            data_in = [deepcopy(subject) for _ in range(self.aleatoric_realisations)]

            # Create the associated Dataset with the given transforms
            ds = Dataset(data_in, self.transform)

            # Create the DataLoader with the given batch size while taking care for all data to be same size
            #   (pad_list_data_collate)
            dl = DataLoader(
                ds,
                num_workers=0,
                # We internally set batch_size to 1 to avoid padding issues with memory since we are using a lot of
                #   memory to create the multiple realisations of the data based on number of augmentation realisations
                #   and number of examples dropout realisations
                batch_size=1,
                collate_fn=pad_list_data_collate,
                pin_memory=True,
                drop_last=False,
                shuffle=False,
                persistent_workers=False,
            )

            # Get the transform_key that will be used to trace the applied transformations of the original data.
            #   Internally it is composed of str(self.orig_key) + TraceKeys.KEY_SUFFIX where
            #   TraceKeys.KEY_SUFFIX = "_transforms"
            transform_key = InvertibleTransform.trace_key(self.orig_key)

            # Automatically inverse the transformations of a given batch of data
            inverter = BatchInverseTransform(
                self.transform, dl, collate_fn=list_data_collate
            )

            outputs: List[np.ndarray] = []

            # Iterate over the data loader and print progress if required
            for batch_data in tqdm(dl) if has_tqdm and self.progress else dl:

                # Grab the data based on given key and move it to device
                batch_images = batch_data[self.image_key].to(self.device)

                # For each epistemic realisation (if self.epistemic_realisations > 1)
                for _ in range(self.epistemic_realisations):
                    gc.collect()
                    # Do the forward pass and make sure the output is a torch.Tensor on CPU
                    batch_output = self.inferer(
                        inputs=batch_images, network=self.network
                    )

                    batch_output[0, :, :, :] = self._post_inferer_transform(
                        batch_output[0, :, :, :]
                    )

                    if isinstance(batch_output, torch.Tensor):
                        batch_output = batch_output.detach().cpu()
                    if isinstance(batch_output, np.ndarray):
                        batch_output = torch.Tensor(batch_output)

                    # If we have aleatoric realisations
                    if self.aleatoric_realisations > 1:
                        # Get the transformations info
                        transform_info = batch_data.get(transform_key, None)

                        if transform_info is None:
                            # In case no invertible transforms, adding dummy info for identity invertible
                            transform_info = [
                                [TraceKeys.NONE] for _ in range(self.batch_size)
                            ]

                        if self.nearest_interp:
                            # Change the interpolation to nearest neighbor for the inverse transform this is done sine we are
                            #   dealing with discrete labels of the segmentation
                            transform_info = convert_inverse_interp_mode(
                                trans_info=deepcopy(transform_info),
                                mode="nearest",
                                align_corners=None,
                            )

                        # Create a dictionary containing the inferred batch and their transforms
                        inferred_dict = {
                            self.orig_key: batch_output,
                            transform_key: transform_info,
                        }

                        # If meta dict is present, add that too (required for some inverse transforms)
                        meta_dict_key = (
                            self.orig_meta_keys
                            or f"{self.orig_key}_{self.meta_key_postfix}"
                        )
                        if meta_dict_key in batch_data:
                            inferred_dict[meta_dict_key] = batch_data[meta_dict_key]

                        # Do inverse transformation (allow missing keys as only inverting the orig_key)
                        with allow_missing_keys_mode(self.transform):  # type: ignore
                            inv_batch = inverter(inferred_dict)

                        batch_output = inv_batch[self.orig_key]

                    # Append output to list
                    outputs.append(batch_output)

            outputs: np.ndarray = np.concatenate(outputs)
            gc.collect()

            if self.return_full_data:
                return outputs

            # calculate metrics
            outputs, *_ = convert_data_type(
                outputs, output_type=torch.Tensor, dtype=torch.float16
            )
            gc.collect()

            # Putting the results into the batch_data where we have the meta information required by the
            #   self._output_trans such as file name
            batch_data["mean"] = torch.mean(outputs, dim=0, keepdim=True)
            # We are only saving the actual mode not the indices
            batch_data["mode"] = torch.mode(outputs, dim=0, keepdim=True)[0]
            batch_data["std"] = torch.std(outputs, dim=0, keepdim=True)

            batch_data = self._decolate(batch_data)

            self._output_transform(batch_data)