From 0462fdb3599fc0597191154c75383bdc42337659 Mon Sep 17 00:00:00 2001 From: IvanTomilov1 Date: Mon, 8 Jan 2024 03:21:23 +0300 Subject: [PATCH] Gaussian wrapper added --- eXNN/bayes/__init__.py | 2 +- eXNN/bayes/api.py | 32 +++++++++++++++ eXNN/bayes/wrapper.py | 92 ++++++++++++++++++++++++++++++++---------- tests/tests.py | 4 +- 4 files changed, 107 insertions(+), 23 deletions(-) diff --git a/eXNN/bayes/__init__.py b/eXNN/bayes/__init__.py index 1b308f6..030f688 100644 --- a/eXNN/bayes/__init__.py +++ b/eXNN/bayes/__init__.py @@ -1 +1 @@ -from .api import DropoutBayesianWrapper +from .api import DropoutBayesianWrapper, DropoutGaussianWrapper diff --git a/eXNN/bayes/api.py b/eXNN/bayes/api.py index 62a2f73..f53845a 100755 --- a/eXNN/bayes/api.py +++ b/eXNN/bayes/api.py @@ -45,3 +45,35 @@ def predict(self, data, n_iter) -> Dict[str, torch.Tensor]: """ res = self.model.mean_forward(data, n_iter) return {"mean": res[0], "std": res[1]} + + +class DropoutGaussianWrapper: + def __init__( + self, + model: torch.nn.Module, + sigma: float + ): + """Class representing bayesian equivalent of a neural network. + + Args: + model (torch.nn.Module): neural network + sigma (float): std of parameters gaussian noise + """ + self.model = create_dropout_bayesian_wrapper(model, "gaussian", sigma = sigma) + + def predict(self, data, n_iter) -> Dict[str, torch.Tensor]: + """Function computes mean and standard deviation of bayesian equivalent + of a neural network. + + Args: + data (_type_): input data of shape NxC1x...xCk, + where N is the number of data points, + C1,...,Ck are dimensions of each data point + n_iter (_type_): number of samplings form the bayesian equivalent + of a neural network + + Returns: + Dict[str, torch.Tensor]: dictionary with `mean` and `std` of prediction + """ + res = self.model.mean_forward(data, n_iter) + return {"mean": res[0], "std": res[1]} diff --git a/eXNN/bayes/wrapper.py b/eXNN/bayes/wrapper.py index ed69b2e..9d83d45 100644 --- a/eXNN/bayes/wrapper.py +++ b/eXNN/bayes/wrapper.py @@ -13,40 +13,51 @@ def __init__( layer: nn.Module, p: Optional[float] = None, a: Optional[float] = None, - b: Optional[float] = None + b: Optional[float] = None, + sigma: Optional[float] = None ): super(ModuleBayesianWrapper, self).__init__() - pab_check = "You can either specify p (simple dropout), or specify a and b (beta dropout)" - assert (p is not None) != ((a is not None) and (b is not None)), pab_check + # Variables correctness checks + pab_check = "You can either specify the following options (exclusively):\n - p (simple dropout)\n - a and b (beta dropout)\n - sigma (gaussian dropout)" + assert (p is not None and a is None and b is None and sigma is None) or \ + (p is None and a is not None and b is not None and sigma is None) or \ + (p is None and a is None and b is None and sigma is not None), pab_check + + if (p is None) and (sigma is None): + ab_check = "If you choose to specify a and b, you must to specify both" + assert (self.a is not None) and (self.b is not None), ab_check - if not type(layer) in [nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d]: - # At the moment we are only modifying linear and convolutional layers - self.layer = layer - else: - self.layer = layer + # At the moment we are only modifying linear and convolutional layers, so check this + self.layer = layer - self.p = p - self.a, self.b = a, b + if type(layer) in [nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d]: + self.p, self.a, self.b, self.sigma = p, a, b, sigma - if self.p is None: - ab_check = "If you choose to specify a and b, you must to specify both" - assert (self.a is not None) and (self.b is not None), ab_check - def dropout_weights(self, weights, bias): - if self.p is not None: - p = self.p - else: - p = Beta(torch.tensor(self.a), torch.tensor(self.b)).sample() + def augment_weights(self, weights, bias): - weights = F.dropout(weights, p, training=True) - bias = F.dropout(bias, p, training=True) + # Check if dropout is chosen + if (self.p is not None) or (self.a is not None and self.b is not None): + # Select correct option and apply dropout + if self.p is not None: + p = self.p + else: + p = Beta(torch.tensor(self.a), torch.tensor(self.b)).sample() + + weights = F.dropout(weights, p, training=True) + bias = F.dropout(bias, p, training=True) + + else: + # If gauss is chosen, then apply it + weights = weights + (torch.randn(*weights.shape)*self.sigma).to(weights.device()) + bias = bias + (torch.randn(*bias.shape)*self.sigma).to(bias.device()) return weights, bias def forward(self, x): - weight, bias = self.dropout_weights(self.layer.weight, self.layer.bias) + weight, bias = self.augment_weights(self.layer.weight, self.layer.bias) if isinstance(self.layer, nn.Linear): return F.linear(x, weight, bias) @@ -141,6 +152,41 @@ def mean_forward( dim=0, ) return results + + +class NetworkBayesGauss(nn.Module): + def __init__( + self, + model: torch.nn.Module, + sigma: float + ): + + super(NetworkBayesGauss, self).__init__() + self.model = copy.deepcopy(model) + self.model = replace_modules_with_wrapper(self.model, + ModuleBayesianWrapper, + {"sigma": sigma}) + + def mean_forward( + self, + data: torch.Tensor, + n_iter: int, + ): + + results = [] + for _ in range(n_iter): + results.append(self.model.forward(data)) + + results = torch.stack(results, dim=1) + + results = torch.stack( + [ + torch.mean(results, dim=1), + torch.std(results, dim=1), + ], + dim=0, + ) + return results def create_dropout_bayesian_wrapper( @@ -149,6 +195,7 @@ def create_dropout_bayesian_wrapper( p: Optional[float] = None, a: Optional[float] = None, b: Optional[float] = None, + sigma: Optional[float] = None ) -> torch.nn.Module: if mode == "basic": net = NetworkBayes(model, p) @@ -156,4 +203,7 @@ def create_dropout_bayesian_wrapper( elif mode == "beta": net = NetworkBayesBeta(model, a, b) + elif mode == 'gauss': + net = NetworkBayesGauss(model, sigma) + return net diff --git a/tests/tests.py b/tests/tests.py index 295fd3d..ef689f8 100644 --- a/tests/tests.py +++ b/tests/tests.py @@ -52,7 +52,7 @@ def test_visualization(): def _test_bayes_prediction(mode: str): - params = {"basic": dict(mode="basic", p=0.5), "beta": dict(mode="beta", a=0.9, b=0.2)} + params = {"basic": dict(mode="basic", p=0.5), "beta": dict(mode="beta", a=0.9, b=0.2), "gauss": dict(mode="gauss", sigma = 1e-2)} N, dim, data = utils.create_testing_data() model = utils.create_testing_model() @@ -73,6 +73,8 @@ def test_basic_bayes_wrapper(): def test_beta_bayes_wrapper(): _test_bayes_prediction("beta") +def test_gauss_bayes_wrapper(): + _test_bayes_prediction("gauss") def test_data_barcode(): N, dim, data = utils.create_testing_data()