Module fl_server_ai.uncertainty.mc_dropout¶

View Source

# SPDX-FileCopyrightText: 2024 Benedikt Franke <benedikt.franke@dlr.de>
# SPDX-FileCopyrightText: 2024 Florian Heinrich <florian.heinrich@dlr.de>
#
# SPDX-License-Identifier: Apache-2.0

import torch
from torch import Tensor
from torch.nn import Module
from torch.nn.modules.dropout import _DropoutNd
from typing import Any, Dict, Tuple

from fl_server_core.models import Model
from fl_server_core.utils.torch_serialization import is_torchscript_instance

from .base import UncertaintyBase


def set_dropout(model: Module, state: bool = True):
    """
    Set the state of the dropout layers to enable or disable them even during inference.

    Args:
        model (Module): PyTorch module
        state (bool, optional): Enable or disable dropout layers. Defaults to True.
    """
    is_torchscript_model = is_torchscript_instance(model)
    for m in model.modules():
        name = m.original_name if is_torchscript_model else m.__class__.__name__
        if isinstance(m, _DropoutNd) or name.lower().__contains__("dropout"):
            m.train(mode=state)


class MCDropout(UncertaintyBase):
    """
    Monte Carlo (MC) Dropout Uncertainty Estimation

    Requirements:

    - model with dropout layers
    - T, number of samples per input (number of monte-carlo samples/forward passes)

    References:

    - Paper: Understanding Measures of Uncertainty for Adversarial Example Detection
             <https://arxiv.org/abs/1803.08533>
    - Code inspiration: <https://github.com/lsgos/uncertainty-adversarial-paper/tree/master>
    """

    @classmethod
    def prediction(cls, input: Tensor, model: Model) -> Tuple[torch.Tensor, Dict[str, Any]]:
        options = cls.get_options(model)
        N = options.get("N", 10)
        softmax = options.get("softmax", False)

        net: Module = model.get_torch_model()
        net.eval()
        set_dropout(net, state=True)

        out_list = []
        for _ in range(N):
            output = net(input).detach()
            # convert to probabilities if necessary
            if softmax:
                output = torch.softmax(output, dim=1)
            out_list.append(output)
        out = torch.stack(out_list, dim=0)  # (n_mc, batch_size, n_classes)

        inference = out.mean(dim=0)
        uncertainty = cls.interpret(out)
        return inference, uncertainty

Functions¶

set_dropout¶

def set_dropout(
    model: torch.nn.modules.module.Module,
    state: bool = True
)

Set the state of the dropout layers to enable or disable them even during inference.

Parameters:

Name	Type	Description	Default
model	Module	PyTorch module	None
state	bool	Enable or disable dropout layers. Defaults to True.	True

View Source

def set_dropout(model: Module, state: bool = True):
    """
    Set the state of the dropout layers to enable or disable them even during inference.

    Args:
        model (Module): PyTorch module
        state (bool, optional): Enable or disable dropout layers. Defaults to True.
    """
    is_torchscript_model = is_torchscript_instance(model)
    for m in model.modules():
        name = m.original_name if is_torchscript_model else m.__class__.__name__
        if isinstance(m, _DropoutNd) or name.lower().__contains__("dropout"):
            m.train(mode=state)

Classes¶

MCDropout¶

class MCDropout(
    /,
    *args,
    **kwargs
)

Monte Carlo (MC) Dropout Uncertainty Estimation

Requirements:

model with dropout layers
T, number of samples per input (number of monte-carlo samples/forward passes)

References:

Paper: Understanding Measures of Uncertainty for Adversarial Example Detection https://arxiv.org/abs/1803.08533
Code inspiration: https://github.com/lsgos/uncertainty-adversarial-paper/tree/master

View Source

class MCDropout(UncertaintyBase):
    """
    Monte Carlo (MC) Dropout Uncertainty Estimation

    Requirements:

    - model with dropout layers
    - T, number of samples per input (number of monte-carlo samples/forward passes)

    References:

    - Paper: Understanding Measures of Uncertainty for Adversarial Example Detection
             <https://arxiv.org/abs/1803.08533>
    - Code inspiration: <https://github.com/lsgos/uncertainty-adversarial-paper/tree/master>
    """

    @classmethod
    def prediction(cls, input: Tensor, model: Model) -> Tuple[torch.Tensor, Dict[str, Any]]:
        options = cls.get_options(model)
        N = options.get("N", 10)
        softmax = options.get("softmax", False)

        net: Module = model.get_torch_model()
        net.eval()
        set_dropout(net, state=True)

        out_list = []
        for _ in range(N):
            output = net(input).detach()
            # convert to probabilities if necessary
            if softmax:
                output = torch.softmax(output, dim=1)
            out_list.append(output)
        out = torch.stack(out_list, dim=0)  # (n_mc, batch_size, n_classes)

        inference = out.mean(dim=0)
        uncertainty = cls.interpret(out)
        return inference, uncertainty

Ancestors (in MRO)¶

fl_server_ai.uncertainty.base.UncertaintyBase
abc.ABC

Static methods¶

expected_entropy¶

def expected_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the mean entropy of the predictive distribution across the MC samples.

Parameters:

Name	Type	Description	Default
predictions	torch.Tensor	predictions of shape (n_mc x batch_size x n_classes)	None

Returns:

Type	Description
torch.Tensor	mean entropy of the predictive distribution

View Source

    @classmethod
    def expected_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the mean entropy of the predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: mean entropy of the predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions).entropy().mean(dim=0)

get_options¶

def get_options(
    obj: fl_server_core.models.model.Model | fl_server_core.models.training.Training
) -> Dict[str, Any]

Get uncertainty options from training options.

Parameters:

Name	Type	Description	Default
obj	Model	Training	The Model or Training object to retrieve options for.

Returns:

Type	Description
Dict[str, Any]	Uncertainty options.

Raises:

Type	Description
TypeError	If the given object is not a Model or Training.

View Source

    @classmethod
    def get_options(cls, obj: Model | Training) -> Dict[str, Any]:
        """
        Get uncertainty options from training options.

        Args:
            obj (Model | Training): The Model or Training object to retrieve options for.

        Returns:
            Dict[str, Any]: Uncertainty options.

        Raises:
            TypeError: If the given object is not a Model or Training.
        """
        if isinstance(obj, Model):
            return Training.objects.filter(model=obj) \
                .values("options") \
                .first()["options"] \
                .get("uncertainty", {})
        if isinstance(obj, Training):
            return obj.options.get("uncertainty", {})
        raise TypeError(f"Expected Model or Training, got {type(obj)}")

interpret¶

def interpret(
    outputs: torch.Tensor
) -> Dict[str, Any]

Interpret the different network (model) outputs and calculate the uncertainty.

Parameters:

Name	Type	Description	Default
outputs	torch.Tensor	multiple network (model) outputs (N, batch_size, n_classes)	None

View Source

    @classmethod
    def interpret(cls, outputs: torch.Tensor) -> Dict[str, Any]:
        """
        Interpret the different network (model) outputs and calculate the uncertainty.

        Args:
            outputs (torch.Tensor): multiple network (model) outputs (N, batch_size, n_classes)

        Return:
            Tuple[torch.Tensor, Dict[str, Any]]: inference and uncertainty
        """
        variance = outputs.var(dim=0)
        std = outputs.std(dim=0)
        if not (torch.all(outputs <= 1.) and torch.all(outputs >= 0.)):
            return dict(variance=variance, std=std)

        predictive_entropy = cls.predictive_entropy(outputs)
        expected_entropy = cls.expected_entropy(outputs)
        mutual_info = predictive_entropy - expected_entropy  # see cls.mutual_information
        return dict(
            variance=variance,
            std=std,
            predictive_entropy=predictive_entropy,
            expected_entropy=expected_entropy,
            mutual_info=mutual_info,
        )

mutual_information¶

def mutual_information(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;

the difference between the mean of the entropy and the entropy of the mean of the predicted distribution on the predictions. This method is also sometimes referred to as the mutual information (MI).

Parameters:

Name	Type	Description	Default
predictions	torch.Tensor	predictions of shape (n_mc x batch_size x n_classes)	None

Returns:

Type	Description
torch.Tensor	difference between the mean of the entropy and the entropy of the mean of the predicted distribution

View Source

    @classmethod
    def mutual_information(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the BALD (Bayesian Active Learning by Disagreement) of a model;
        the difference between the mean of the entropy and the entropy of the mean
        of the predicted distribution on the predictions.
        This method is also sometimes referred to as the mutual information (MI).

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: difference between the mean of the entropy and the entropy of the mean
                    of the predicted distribution
        """
        return cls.predictive_entropy(predictions) - cls.expected_entropy(predictions)

prediction¶

def prediction(
    input: torch.Tensor,
    model: fl_server_core.models.model.Model
) -> Tuple[torch.Tensor, Dict[str, Any]]

Make a prediction using the given input and model.

Parameters:

Name	Type	Description	Default
input	torch.Tensor	The input to make a prediction for.	None
model	Model	The model to use for making the prediction.	None

Returns:

Type	Description
Tuple[torch.Tensor, Dict[str, Any]]	The prediction and any associated uncertainty.

View Source

    @classmethod
    def prediction(cls, input: Tensor, model: Model) -> Tuple[torch.Tensor, Dict[str, Any]]:
        options = cls.get_options(model)
        N = options.get("N", 10)
        softmax = options.get("softmax", False)

        net: Module = model.get_torch_model()
        net.eval()
        set_dropout(net, state=True)

        out_list = []
        for _ in range(N):
            output = net(input).detach()
            # convert to probabilities if necessary
            if softmax:
                output = torch.softmax(output, dim=1)
            out_list.append(output)
        out = torch.stack(out_list, dim=0)  # (n_mc, batch_size, n_classes)

        inference = out.mean(dim=0)
        uncertainty = cls.interpret(out)
        return inference, uncertainty

predictive_entropy¶

def predictive_entropy(
    predictions: torch.Tensor
) -> torch.Tensor

Calculate the entropy of the mean predictive distribution across the MC samples.

Parameters:

Name	Type	Description	Default
predictions	torch.Tensor	predictions of shape (n_mc x batch_size x n_classes)	None

Returns:

Type	Description
torch.Tensor	entropy of the mean predictive distribution

View Source

    @classmethod
    def predictive_entropy(cls, predictions: torch.Tensor) -> torch.Tensor:
        """
        Calculate the entropy of the mean predictive distribution across the MC samples.

        Args:
            predictions (torch.Tensor): predictions of shape (n_mc x batch_size x n_classes)

        Returns:
            torch.Tensor: entropy of the mean predictive distribution
        """
        return torch.distributions.Categorical(probs=predictions.mean(dim=0)).entropy()

to_json¶

def to_json(
    inference: torch.Tensor,
    uncertainty: Dict[str, Any] = {},
    **json_kwargs
) -> str

Convert the given inference and uncertainty data to a JSON string.

Parameters:

Name	Type	Description	Default
inference	torch.Tensor	The inference to convert.	None
uncertainty	Dict[str, Any]	The uncertainty to convert.	None
**json_kwargs	None	Additional keyword arguments to pass to `json.dumps`.	None

Returns:

Type	Description
str	A JSON string representation of the given inference and uncertainty data.

View Source

    @classmethod
    def to_json(cls, inference: torch.Tensor, uncertainty: Dict[str, Any] = {}, **json_kwargs) -> str:
        """
        Convert the given inference and uncertainty data to a JSON string.

        Args:
            inference (torch.Tensor): The inference to convert.
            uncertainty (Dict[str, Any]): The uncertainty to convert.
            **json_kwargs: Additional keyword arguments to pass to `json.dumps`.

        Returns:
            str: A JSON string representation of the given inference and uncertainty data.
        """
        def prepare(v):
            if isinstance(v, torch.Tensor):
                return v.cpu().tolist()
            if isinstance(v, np.ndarray):  # cspell:ignore ndarray
                return v.tolist()
            if isinstance(v, dict):
                return {k: prepare(_v) for k, _v in v.items()}
            return v

        return json.dumps({
            "inference": inference.tolist(),
            "uncertainty": prepare(uncertainty) if uncertainty else {},
        }, **json_kwargs)