eXNN.NetBayesianization.api

 1from typing import Dict, Optional
 2import torch
 3import torch.optim
 4from eXNN.NetBayesianization.wrap import create_bayesian_wrapper
 5
 6
 7def BasicBayesianPrediction(data: torch.Tensor,
 8                            model: torch.nn.Module,
 9                            n_iter: int,
10                            mode: str,
11                            p: Optional[float] = None,
12                            a: Optional[float] = None,
13                            b: Optional[float] = None) -> Dict[str, torch.Tensor]:
14    """Function computes mean and standard deviation of bayesian equivalent
15        of a neural network.
16
17    Args:
18        data (torch.Tensor): input data of shape NxC1x...xCk,
19            where N is the number of data points,
20            C1,...,Ck are dimensions of each data point
21        model (torch.nn.Module): neural network
22        n_iter (int): number of samplings form the bayesian equivalent of a neural network
23        mode (str): bayesianization method (`basic` or `beta`)
24        p (Optional[float], optional): parameter of dropout (for `basic`
25            bayesianization). Defaults to None.
26        a (Optional[float], optional): parameter of beta distribution (for `beta`
27            bayesianization). Defaults to None.
28        b (Optional[float], optional): parameter of beta distribution (for `beta`
29            bayesianization). Defaults to None.
30
31    Returns:
32        Dict[str, torch.Tensor]: dictionary with `mean` and `std` of prediction
33    """
34    return BasicBayesianWrapper(model, mode, p=p, a=a, b=b).predict(data, n_iter)
35
36
37class BasicBayesianWrapper:
38    def __init__(self,
39                 model: torch.nn.Module,
40                 mode: str,
41                 p: Optional[float] = None,
42                 a: Optional[float] = None,
43                 b: Optional[float] = None):
44        """Class representing bayesian equivalent of a neural network.
45
46        Args:
47            model (torch.nn.Module): neural network
48            mode (str): bayesianization method (`basic` or `beta`)
49            p (Optional[float], optional): parameter of dropout (for `basic`
50                bayesianization). Defaults to None.
51            a (Optional[float], optional): parameter of beta distribution (for `beta`
52                bayesianization). Defaults to None.
53            b (Optional[float], optional): parameter of beta distribution (for `beta`
54                bayesianization). Defaults to None.
55        """
56        self.model = create_bayesian_wrapper(model, mode, p=p, a=a, b=b)
57
58    def predict(self, data, n_iter) -> Dict[str, torch.Tensor]:
59        """Function computes mean and standard deviation of bayesian equivalent
60            of a neural network.
61
62        Args:
63            data (_type_): input data of shape NxC1x...xCk,
64                where N is the number of data points,
65                C1,...,Ck are dimensions of each data point
66            n_iter (_type_): number of samplings form the bayesian equivalent
67                of a neural network
68
69        Returns:
70            Dict[str, torch.Tensor]: dictionary with `mean` and `std` of prediction
71        """
72        res = self.model.mean_forward(data, n_iter)
73        return {'mean': res[0], 'std': res[1]}