core.layer.AbstractProbLayer

 1from abc import ABC
 2
 3from torch import Tensor, nn
 4
 5from core.distribution import AbstractVariable
 6
 7
 8class AbstractProbLayer(nn.Module, ABC):
 9    """
10    Base class for probabilistic layers in a neural network.
11
12    It introduces the concept of `probabilistic_mode`, which determines whether
13    a layer samples parameters from a distribution or uses deterministic values.
14    Each layer holds references to distributions for weights/biases and their prior.
15    """
16
17    probabilistic_mode: bool
18    _weight_dist: AbstractVariable
19    _bias_dist: AbstractVariable
20    _prior_weight_dist: AbstractVariable
21    _prior_bias_dist: AbstractVariable
22
23    def probabilistic(self, mode: bool = True):
24        """
25        Recursively set the `probabilistic_mode` flag for this layer and its children.
26
27        Args:
28            mode (bool, optional): If True, the layer will draw samples from its
29                distributions during forward passes. If False, it will typically
30                use deterministic parameters (i.e., the mean), except in training.
31        """
32        if not isinstance(mode, bool):
33            raise ValueError("probabilistic mode is expected to be boolean")
34        if isinstance(self, AbstractProbLayer):
35            self.probabilistic_mode = mode
36        for module in self.children():
37            if hasattr(module, "probabilistic_mode") and hasattr(
38                module, "probabilistic"
39            ):
40                module.probabilistic(mode)
41
42    def sample_from_distribution(self) -> tuple[Tensor, Tensor]:
43        """
44        Draw samples for weight and bias from their respective distributions.
45
46        Behavior depends on `probabilistic_mode`:
47          - If `probabilistic_mode` is True, the method samples from `_weight_dist`
48            and `_bias_dist`.
49          - If `probabilistic_mode` is False and the layer is in eval mode, it uses
50            the mean (mu) of each distribution.
51          - If `probabilistic_mode` is False and the layer is in training mode, it
52            raises an error, as training requires sampling.
53
54        Returns:
55            Tuple[Tensor, Tensor]: Sampled weight and bias. If the layer has no bias,
56            the second returned value is None.
57        """
58        if self.probabilistic_mode:
59            sampled_weight = self._weight_dist.sample()
60            sampled_bias = self._bias_dist.sample() if self._bias_dist else None
61        else:
62            if not self.training:
63                sampled_weight = self._weight_dist.mu
64                sampled_bias = self._bias_dist.mu if self._bias_dist else None
65            else:
66                raise ValueError("Only training with probabilistic mode is allowed")
67        return sampled_weight, sampled_bias