Skip to content

Base Network Distributions

BaseNetDistribution

Base class for distribution of nets. This class sees nets as elements of distribution. It helps sample nets from this distribution or estimate statistics of distribution. For this purpose it have base module architecture and distribution of parameters for each of it weights.

Source code in src/methods/bayes/base/net_distribution.py 
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
class BaseNetDistribution:
    """
    Base class for distribution of nets. This class sees nets as elements of distribution.
    It helps sample nets from this distribution or estimate statistics of distribution.
    For this purpose it have base module architecture and distribution of parameters for
    each of it weights.
    """

    def __init__(
        self, base_module: nn.Module, weight_distribution: dict[str, ParamDist]
    ) -> None:
        """_summary_

        Args:
            base_module (nn.Module): custom module layer which is going to be converted to BayesModule
            weight_distribution (dict[str, ParamDist]): posteror distribution for each parameter of moudule
        """
        super().__init__()
        self.base_module: nn.Module = base_module
        """Show default architecture of module for which should evalute parameters"""
        self.weight_distribution: dict[str, ParamDist] = weight_distribution
        """Distribution of parameter for each named parameter of base_module"""

    def detach_(self):
        """
        Detach(Made deepcopy) base module from original module
        """
        self.base_module = copy.deepcopy(self.base_module)

    def sample_params(self) -> dict[str, nn.Parameter]:
        """
        Sample only model parameter from distribution and
        return it.

        Returns:
            dict[str, nn.Parameter]: Return dict of sampled parameters, where
                key is name of parameter, value is valeu of parameter
        """
        param_sample_dict: dict[str, nn.Parameter] = {}
        for param_name, param_posterior in self.weight_distribution.items():
            param_sample = param_posterior.rsample()
            param_sample_dict[param_name] = nn.Parameter(param_sample)
            del_attr(self.base_module, param_name.split("."))
            set_attr(self.base_module, param_name.split("."), param_sample)
        return param_sample_dict

    def sample_model(self) -> nn.Module:
        """
        Sample only model from distribution and
        return it. Note that model is the same that in base_module.

        Returns:
            nn.Module: sampled base_module with sampled parameters
        """
        self.sample_params()
        return self.base_module

    def set_map_params(self) -> None:
        """
        Set MAP estimation parameters of model from distribution and
        set it. So it sets most probable model.
        """
        for param_name, dist in self.weight_distribution.items():
            pt = dist.map
            # pt = torch.nn.Parameter(pt.to_sparse())
            set_attr(self.base_module, param_name.split("."), pt)

    def set_mean_params(self) -> None:
        """
        Set unbaised estimation parameters of model from distribution and
        set it. So it sets model that would be returned at means.
        """
        for param_name, dist in self.weight_distribution.items():
            pt = dist.mean
            # pt = torch.nn.Parameter(pt.to_sparse())
            set_attr(self.base_module, param_name.split("."), pt)

    def __replace_with_parameters(self) -> None:
        """
        Replace parameters of base_module wit nn.Parameters
        """
        for param_name in self.weight_distribution.keys():
            pt = get_attr(self.base_module, param_name.split("."))
            pt = nn.Parameter(pt)
            set_attr(self.base_module, param_name.split("."), pt)

    def get_model(self) -> nn.Module:
        """
        Get model with last set parameters. (It still we be the same model as in base_module)
        """
        self.__replace_with_parameters()
        return self.base_module

    def get_model_snapshot(self) -> nn.Module:
        """
        Get deepcopy of model with last set parameters.
        """
        return copy.deepcopy(self.get_model())

base_module: nn.Module = base_module instance-attribute

Show default architecture of module for which should evalute parameters

weight_distribution: dict[str, ParamDist] = weight_distribution instance-attribute

Distribution of parameter for each named parameter of base_module

__init__(base_module, weight_distribution)

summary

Parameters:

Name Type Description Default
base_module Module

custom module layer which is going to be converted to BayesModule

 required
weight_distribution dict[str, ParamDist]

posteror distribution for each parameter of moudule

 required
Source code in src/methods/bayes/base/net_distribution.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def __init__(
    self, base_module: nn.Module, weight_distribution: dict[str, ParamDist]
) -> None:
    """_summary_

    Args:
        base_module (nn.Module): custom module layer which is going to be converted to BayesModule
        weight_distribution (dict[str, ParamDist]): posteror distribution for each parameter of moudule
    """
    super().__init__()
    self.base_module: nn.Module = base_module
    """Show default architecture of module for which should evalute parameters"""
    self.weight_distribution: dict[str, ParamDist] = weight_distribution
    """Distribution of parameter for each named parameter of base_module"""

__replace_with_parameters()

Replace parameters of base_module wit nn.Parameters

Source code in src/methods/bayes/base/net_distribution.py 
86
87
88
89
90
91
92
93
def __replace_with_parameters(self) -> None:
    """
    Replace parameters of base_module wit nn.Parameters
    """
    for param_name in self.weight_distribution.keys():
        pt = get_attr(self.base_module, param_name.split("."))
        pt = nn.Parameter(pt)
        set_attr(self.base_module, param_name.split("."), pt)

detach_()

Detach(Made deepcopy) base module from original module

Source code in src/methods/bayes/base/net_distribution.py 
32
33
34
35
36
def detach_(self):
    """
    Detach(Made deepcopy) base module from original module
    """
    self.base_module = copy.deepcopy(self.base_module)

get_model()

Get model with last set parameters. (It still we be the same model as in base_module)

Source code in src/methods/bayes/base/net_distribution.py 
 95
 96
 97
 98
 99
100
def get_model(self) -> nn.Module:
    """
    Get model with last set parameters. (It still we be the same model as in base_module)
    """
    self.__replace_with_parameters()
    return self.base_module

get_model_snapshot()

Get deepcopy of model with last set parameters.

Source code in src/methods/bayes/base/net_distribution.py 
102
103
104
105
106
def get_model_snapshot(self) -> nn.Module:
    """
    Get deepcopy of model with last set parameters.
    """
    return copy.deepcopy(self.get_model())

sample_model()

Sample only model from distribution and return it. Note that model is the same that in base_module.

Returns:

Type Description
Module

nn.Module: sampled base_module with sampled parameters
Source code in src/methods/bayes/base/net_distribution.py 
55
56
57
58
59
60
61
62
63
64
def sample_model(self) -> nn.Module:
    """
    Sample only model from distribution and
    return it. Note that model is the same that in base_module.

    Returns:
        nn.Module: sampled base_module with sampled parameters
    """
    self.sample_params()
    return self.base_module

sample_params()

Sample only model parameter from distribution and return it.

Returns:

Type Description
dict[str, Parameter]

dict[str, nn.Parameter]: Return dict of sampled parameters, where key is name of parameter, value is valeu of parameter
Source code in src/methods/bayes/base/net_distribution.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def sample_params(self) -> dict[str, nn.Parameter]:
    """
    Sample only model parameter from distribution and
    return it.

    Returns:
        dict[str, nn.Parameter]: Return dict of sampled parameters, where
            key is name of parameter, value is valeu of parameter
    """
    param_sample_dict: dict[str, nn.Parameter] = {}
    for param_name, param_posterior in self.weight_distribution.items():
        param_sample = param_posterior.rsample()
        param_sample_dict[param_name] = nn.Parameter(param_sample)
        del_attr(self.base_module, param_name.split("."))
        set_attr(self.base_module, param_name.split("."), param_sample)
    return param_sample_dict

set_map_params()

Set MAP estimation parameters of model from distribution and set it. So it sets most probable model.

Source code in src/methods/bayes/base/net_distribution.py 
66
67
68
69
70
71
72
73
74
def set_map_params(self) -> None:
    """
    Set MAP estimation parameters of model from distribution and
    set it. So it sets most probable model.
    """
    for param_name, dist in self.weight_distribution.items():
        pt = dist.map
        # pt = torch.nn.Parameter(pt.to_sparse())
        set_attr(self.base_module, param_name.split("."), pt)

set_mean_params()

Set unbaised estimation parameters of model from distribution and set it. So it sets model that would be returned at means.

Source code in src/methods/bayes/base/net_distribution.py 
76
77
78
79
80
81
82
83
84
def set_mean_params(self) -> None:
    """
    Set unbaised estimation parameters of model from distribution and
    set it. So it sets model that would be returned at means.
    """
    for param_name, dist in self.weight_distribution.items():
        pt = dist.mean
        # pt = torch.nn.Parameter(pt.to_sparse())
        set_attr(self.base_module, param_name.split("."), pt)