mindspore 2.4.0__cp311-cp311-win_amd64.whl

mindspore/nn/probability/bnn_layers/dense_variational.py

@@ -0,0 +1,302 @@
+# Copyright 2020 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""dense_variational"""
+from mindspore.ops import operations as P
+from mindspore import _checkparam as Validator
+from ...cell import Cell
+from ...layer.activation import get_activation
+from ..distribution.normal import Normal
+from .layer_distribution import NormalPrior, normal_post_fn
+from ._util import check_prior, check_posterior
+
+__all__ = ['DenseReparam', 'DenseLocalReparam']
+
+
+class _DenseVariational(Cell):
+    """
+    Base class for all dense variational layers.
+    """
+
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            activation=None,
+            has_bias=True,
+            weight_prior_fn=NormalPrior,
+            weight_posterior_fn=normal_post_fn,
+            bias_prior_fn=NormalPrior,
+            bias_posterior_fn=normal_post_fn):
+        super(_DenseVariational, self).__init__()
+        self.in_channels = Validator.check_positive_int(in_channels)
+        self.out_channels = Validator.check_positive_int(out_channels)
+        self.has_bias = Validator.check_bool(has_bias)
+
+        self.weight_prior = check_prior(weight_prior_fn, "weight_prior_fn")
+        self.weight_posterior = check_posterior(weight_posterior_fn, shape=[self.out_channels, self.in_channels],
+                                                param_name='bnn_weight', arg_name="weight_posterior_fn")
+
+        if self.has_bias:
+            self.bias_prior = check_prior(bias_prior_fn, "bias_prior_fn")
+            self.bias_posterior = check_posterior(bias_posterior_fn, shape=[self.out_channels], param_name='bnn_bias',
+                                                  arg_name="bias_posterior_fn")
+
+        self.activation = activation
+        if not self.activation:
+            self.activation_flag = False
+        else:
+            self.activation_flag = True
+            if isinstance(self.activation, str):
+                self.activation = get_activation(activation)
+            elif isinstance(self.activation, Cell):
+                self.activation = activation
+            else:
+                raise ValueError('The type of `activation` is wrong.')
+
+        self.matmul = P.MatMul(transpose_b=True)
+        self.bias_add = P.BiasAdd()
+        self.sum = P.ReduceSum()
+
+    def construct(self, x):
+        outputs = self._apply_variational_weight(x)
+        if self.has_bias:
+            outputs = self.apply_variational_bias(outputs)
+        if self.activation_flag:
+            outputs = self.activation(outputs)
+        return outputs
+
+    def extend_repr(self):
+        """Display instance object as string."""
+        s = 'in_channels={}, out_channels={}, weight_mean={}, weight_std={}, has_bias={}' \
+            .format(self.in_channels, self.out_channels, self.weight_posterior.mean,
+                    self.weight_posterior.untransformed_std, self.has_bias)
+        if self.has_bias:
+            s += ', bias_mean={}, bias_std={}' \
+                .format(self.bias_posterior.mean, self.bias_posterior.untransformed_std)
+        if self.activation_flag:
+            s += ', activation={}'.format(self.activation)
+        return s
+
+    def apply_variational_bias(self, inputs):
+        """Calculate bias."""
+        bias_posterior_tensor = self.bias_posterior("sample")
+        return self.bias_add(inputs, bias_posterior_tensor)
+
+    def compute_kl_loss(self):
+        """Compute kl loss"""
+        weight_args_list = self.weight_posterior("get_dist_args")
+        weight_type = self.weight_posterior("get_dist_type")
+
+        kl = self.weight_prior("kl_loss", weight_type, *weight_args_list)
+        kl_loss = self.sum(kl)
+        if self.has_bias:
+            bias_args_list = self.bias_posterior("get_dist_args")
+            bias_type = self.bias_posterior("get_dist_type")
+
+            kl = self.bias_prior("kl_loss", bias_type, *bias_args_list)
+            kl = self.sum(kl)
+            kl_loss += kl
+        return kl_loss
+
+
+class DenseReparam(_DenseVariational):
+    r"""
+    Dense variational layers with Reparameterization.
+
+    For more details, refer to the paper `Auto-Encoding Variational Bayes <https://arxiv.org/abs/1312.6114>`_.
+
+    Applies dense-connected layer to the input. This layer implements the operation as:
+
+    .. math::
+        \text{outputs} = \text{activation}(\text{inputs} * \text{weight} + \text{bias}),
+
+    where :math:`\text{activation}` is the activation function passed as the activation
+    argument (if passed in), :math:`\text{activation}` is a weight matrix with the same
+    data type as the inputs created by the layer, :math:`\text{weight}` is a weight
+    matrix sampling from posterior distribution of weight, and :math:`\text{bias}` is a
+    bias vector with the same data type as the inputs created by the layer (only if
+    has_bias is True). The bias vector is sampling from posterior distribution of
+    :math:`\text{bias}`.
+
+    Args:
+        in_channels (int): The number of input channel.
+        out_channels (int): The number of output channel .
+        activation (str, Cell): A regularization function applied to the output of the layer.
+            The type of `activation` can be a string (eg. 'relu') or a Cell (eg. nn.ReLU()).
+            Note that if the type of activation is Cell, it must be instantiated beforehand.
+            Default: ``None`` .
+        has_bias (bool): Specifies whether the layer uses a bias vector. Default: ``False`` .
+        weight_prior_fn (Cell): The prior distribution for weight.
+            It must return a mindspore distribution instance.
+            Default: ``NormalPrior`` . (which creates an instance of standard
+            normal distribution). The current version only supports normal distribution.
+        weight_posterior_fn (function): The posterior distribution for sampling weight.
+            It must be a function handle which returns a mindspore
+            distribution instance. Default: ``normal_post_fn`` .
+            The current version only supports normal distribution.
+        bias_prior_fn (Cell): The prior distribution for bias vector. It must return
+            a mindspore distribution. Default: ``NormalPrior`` (which creates an
+            instance of standard normal distribution). The current version
+            only supports normal distribution.
+        bias_posterior_fn (function): The posterior distribution for sampling bias vector.
+            It must be a function handle which returns a mindspore
+            distribution instance. Default: ``normal_post_fn`` .
+            The current version only supports normal distribution.
+
+    Inputs:
+        - **input** (Tensor) - The shape of the tensor is :math:`(N, in\_channels)`.
+
+    Outputs:
+        Tensor, the shape of the tensor is :math:`(N, out\_channels)`.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore
+        >>> from mindspore import Tensor
+        >>> from mindspore.nn.probability import bnn_layers
+        >>> net = bnn_layers.DenseReparam(3, 4)
+        >>> input = Tensor(np.random.randint(0, 255, [2, 3]), mindspore.float32)
+        >>> output = net(input).shape
+        >>> print(output)
+        (2, 4)
+    """
+
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            activation=None,
+            has_bias=True,
+            weight_prior_fn=NormalPrior,
+            weight_posterior_fn=normal_post_fn,
+            bias_prior_fn=NormalPrior,
+            bias_posterior_fn=normal_post_fn):
+        super(DenseReparam, self).__init__(
+            in_channels,
+            out_channels,
+            activation=activation,
+            has_bias=has_bias,
+            weight_prior_fn=weight_prior_fn,
+            weight_posterior_fn=weight_posterior_fn,
+            bias_prior_fn=bias_prior_fn,
+            bias_posterior_fn=bias_posterior_fn
+        )
+
+    def _apply_variational_weight(self, inputs):
+        """Calculate weight."""
+        weight_posterior_tensor = self.weight_posterior("sample")
+        outputs = self.matmul(inputs, weight_posterior_tensor)
+        return outputs
+
+
+class DenseLocalReparam(_DenseVariational):
+    r"""
+    Dense variational layers with Local Reparameterization.
+
+    For more details, refer to the paper `Variational Dropout and the Local Reparameterization
+    Trick <https://arxiv.org/abs/1506.02557>`_.
+
+    Applies dense-connected layer to the input. This layer implements the operation as:
+
+    .. math::
+        \text{outputs} = \text{activation}(\text{inputs} * \text{weight} + \text{bias}),
+
+    where :math:`\text{activation}` is the activation function passed as the activation
+    argument (if passed in), :math:`\text{activation}` is a weight matrix with the same
+    data type as the inputs created by the layer, :math:`\text{weight}` is a weight
+    matrix sampling from posterior distribution of weight, and :math:`\text{bias}` is a
+    bias vector with the same data type as the inputs created by the layer (only if
+    has_bias is True). The bias vector is sampling from posterior distribution of
+    :math:`\text{bias}`.
+
+    Args:
+        in_channels (int): The number of input channel.
+        out_channels (int): The number of output channel .
+        activation (str, Cell): A regularization function applied to the output of the layer.
+            The type of `activation` can be a string (eg. 'relu') or a Cell (eg. nn.ReLU()).
+            Note that if the type of activation is Cell, it must be instantiated beforehand.
+            Default: ``None`` .
+        has_bias (bool): Specifies whether the layer uses a bias vector. Default: ``False`` .
+        weight_prior_fn (Cell): The prior distribution for weight.
+            It must return a mindspore distribution instance.
+            Default: ``NormalPrior`` . (which creates an instance of standard
+            normal distribution). The current version only supports normal distribution.
+        weight_posterior_fn (function): The posterior distribution for sampling weight.
+            It must be a function handle which returns a mindspore
+            distribution instance. Default: ``normal_post_fn`` .
+            The current version only supports normal distribution.
+        bias_prior_fn (Cell): The prior distribution for bias vector. It must return
+            a mindspore distribution. Default: ``NormalPrior`` (which creates an
+            instance of standard normal distribution). The current version
+            only supports normal distribution.
+        bias_posterior_fn (function): The posterior distribution for sampling bias vector.
+            It must be a function handle which returns a mindspore
+            distribution instance. Default: ``normal_post_fn`` .
+            The current version only supports normal distribution.
+
+    Inputs:
+        - **input** (Tensor) - The shape of the tensor is :math:`(N, in\_channels)`.
+
+    Outputs:
+        Tensor, the shape of the tensor is :math:`(N, out\_channels)`.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore
+        >>> from mindspore import Tensor
+        >>> from mindspore.nn.probability import bnn_layers
+        >>> net = bnn_layers.DenseLocalReparam(3, 4)
+        >>> input = Tensor(np.random.randint(0, 255, [2, 3]), mindspore.float32)
+        >>> output = net(input).shape
+        >>> print(output)
+        (2, 4)
+    """
+
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            activation=None,
+            has_bias=True,
+            weight_prior_fn=NormalPrior,
+            weight_posterior_fn=normal_post_fn,
+            bias_prior_fn=NormalPrior,
+            bias_posterior_fn=normal_post_fn):
+        super(DenseLocalReparam, self).__init__(
+            in_channels,
+            out_channels,
+            activation=activation,
+            has_bias=has_bias,
+            weight_prior_fn=weight_prior_fn,
+            weight_posterior_fn=weight_posterior_fn,
+            bias_prior_fn=bias_prior_fn,
+            bias_posterior_fn=bias_posterior_fn
+        )
+        self.sqrt = P.Sqrt()
+        self.square = P.Square()
+        self.normal = Normal()
+
+    def _apply_variational_weight(self, inputs):
+        """Calculate weight."""
+        mean = self.matmul(inputs, self.weight_posterior("mean"))
+        std = self.sqrt(self.matmul(self.square(inputs), self.square(self.weight_posterior("sd"))))
+        weight_posterior_affine_tensor = self.normal("sample", mean=mean, sd=std)
+        return weight_posterior_affine_tensor

mindspore/nn/probability/bnn_layers/layer_distribution.py

@@ -0,0 +1,123 @@
+# Copyright 2020 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Initialize normal distributions"""
+import numpy as np
+import mindspore.common.dtype as mstype
+from mindspore.common.tensor import Tensor
+from mindspore.common.parameter import Parameter
+from mindspore.ops import operations as P
+from ...cell import Cell
+from ..distribution.normal import Normal
+
+__all__ = ['NormalPrior', 'NormalPosterior']
+
+
+class NormalPrior(Cell):
+    r"""
+    To initialize a normal distribution of mean 0 and standard deviation 0.1.
+
+    Args:
+        dtype (mindspore.dtype): The argument is used to define the data type of the output tensor.
+            Default: ``mindspore.float32`` .
+        mean (int, float): Mean of normal distribution. Default: ``0`` .
+        std (int, float): Standard deviation of normal distribution. Default: ``0.1`` .
+
+    Returns:
+        Cell, a normal distribution.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+    """
+    def __init__(self, dtype=mstype.float32, mean=0, std=0.1):
+        super(NormalPrior, self).__init__()
+        self.normal = Normal(mean, std, dtype=dtype)
+
+    def construct(self, *inputs):
+        return self.normal(*inputs)
+
+
+class NormalPosterior(Cell):
+    r"""
+    Build Normal distributions with trainable parameters.
+
+    Args:
+        name (str): Name prepended to trainable parameter.
+        shape (list, tuple): Shape of the mean and standard deviation.
+        dtype (mindspore.dtype): The argument is used to define the data type of the output tensor.
+            Default: ``mindspore.float32`` .
+        loc_mean (int, float): Mean of distribution to initialize trainable parameters. Default: ``0`` .
+        loc_std (int, float): Standard deviation of distribution to initialize trainable parameters. Default: ``0.1`` .
+        untransformed_scale_mean (int, float): Mean of distribution to initialize trainable parameters.
+            Default: ``-5`` .
+        untransformed_scale_std (int, float): Standard deviation of distribution to initialize trainable parameters.
+            Default: ``0.1`` .
+
+    Returns:
+        Cell, a normal distribution.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+    """
+    def __init__(self,
+                 name,
+                 shape,
+                 dtype=mstype.float32,
+                 loc_mean=0,
+                 loc_std=0.1,
+                 untransformed_scale_mean=-5,
+                 untransformed_scale_std=0.1):
+        super(NormalPosterior, self).__init__()
+        if not isinstance(name, str):
+            raise TypeError('The type of `name` must be `str`')
+
+        if not isinstance(shape, (tuple, list)):
+            raise TypeError('The type of `shape` must be `tuple` or `list`')
+
+        if isinstance(loc_mean, bool) or not isinstance(loc_mean, (int, float)):
+            raise TypeError('The type of `loc_mean` must be `int` or `float`')
+
+        if isinstance(untransformed_scale_mean, bool) or not isinstance(untransformed_scale_mean, (int, float)):
+            raise TypeError('The type of `untransformed_scale_mean` must be `int` or `float`')
+
+        if isinstance(loc_std, bool) or not (isinstance(loc_std, (int, float)) and loc_std >= 0):
+            raise TypeError('The type of `loc_std` must be `int` or `float` and its value must > 0')
+
+        if isinstance(loc_std, bool) or not (isinstance(untransformed_scale_std, (int, float)) and
+                                             untransformed_scale_std >= 0):
+            raise TypeError('The type of `untransformed_scale_std` must be `int` or `float` and '
+                            'its value must > 0')
+
+        self.mean = Parameter(
+            Tensor(np.random.normal(loc_mean, loc_std, shape), dtype=dtype), name=name + '_mean')
+
+        self.untransformed_std = Parameter(
+            Tensor(np.random.normal(untransformed_scale_mean, untransformed_scale_std, shape), dtype=dtype),
+            name=name + '_untransformed_std')
+
+        self.normal = Normal()
+
+    def _std_trans(self, std_pre):
+        """Transform std_pre to prevent its value being zero."""
+        std = 1e-6 + P.Log()(P.Exp()(std_pre) + 1)
+        return std
+
+    def construct(self, *inputs):
+        std = self._std_trans(self.untransformed_std)
+        return self.normal(*inputs, mean=self.mean, sd=std)
+
+
+def normal_post_fn(name, shape):
+    """Provide normal posterior distribution."""
+    return NormalPosterior(name=name, shape=shape)

mindspore/nn/probability/distribution/__init__.py

@@ -0,0 +1,56 @@
+# Copyright 2020 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""
+Distributions are the high-level components used to construct the probabilistic network.
+"""
+
+from .distribution import Distribution
+from .transformed_distribution import TransformedDistribution
+from .bernoulli import Bernoulli
+from .beta import Beta
+from .categorical import Categorical
+from .cauchy import Cauchy
+from .exponential import Exponential
+from .gamma import Gamma
+from .geometric import Geometric
+from .gumbel import Gumbel
+from .logistic import Logistic
+from .log_normal import LogNormal
+from .normal import Normal
+from .poisson import Poisson
+from .uniform import Uniform
+from .half_normal import HalfNormal
+from .laplace import Laplace
+from .student_t import StudentT
+
+__all__ = ['Distribution',
+           'TransformedDistribution',
+           'Bernoulli',
+           'Beta',
+           'Categorical',
+           'Cauchy',
+           'Exponential',
+           'Gamma',
+           'Geometric',
+           'Gumbel',
+           'Logistic',
+           'LogNormal',
+           'Normal',
+           'Poisson',
+           'Uniform',
+           'HalfNormal',
+           'Laplace',
+           'StudentT',
+           ]

mindspore/nn/probability/distribution/_utils/__init__.py

@@ -0,0 +1,34 @@
+# Copyright 2020 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""
+Distribution operation utility functions.
+"""
+from .custom_ops import exp_generic, log_generic, broadcast_to
+from .utils import cast_to_tensor, check_greater_equal_zero, check_greater_zero
+from .utils import check_greater, check_prob, CheckTensor, CheckTuple, set_param_type
+
+__all__ = [
+    'cast_to_tensor',
+    'check_greater',
+    'check_greater_equal_zero',
+    'check_greater_zero',
+    'check_prob',
+    'exp_generic',
+    'log_generic',
+    'broadcast_to',
+    'set_param_type',
+    'CheckTensor',
+    'CheckTuple',
+]

mindspore/nn/probability/distribution/_utils/custom_ops.py

@@ -0,0 +1,96 @@
+# Copyright 2020 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Utility functions to help distribution class."""
+import numpy as np
+from mindspore.ops import operations as P
+from mindspore.ops import functional as F
+from mindspore.ops.operations import _inner_ops as inner
+from mindspore.ops.primitive import constexpr
+from mindspore.common import dtype as mstype
+from .utils import CheckTensor
+
+
+@constexpr(check=False)
+def _check_tensor(x, name):
+    CheckTensor()(x, name)
+    return x
+
+
+def exp_generic(input_x):
+    """
+    Exp op on Ascend doesn't support int types.
+    Fix this with casting the type.
+    """
+    exp = P.Exp()
+    cast = P.Cast()
+    dtype = P.DType()
+    checktype = inner.IsSubClass()
+
+    if not checktype(dtype(input_x), mstype.float_):
+        input_x = cast(input_x, mstype.float32)
+    return exp(input_x)
+
+
+def log_generic(input_x):
+    """
+    Log op on Ascend is calculated as log(abs(x)).
+    Fix this with putting negative values as nan.
+    And log op on Ascend doesn't support int types.
+    Fix this with casting the type.
+    """
+    log = P.Log()
+    less = P.Less()
+    lessequal = P.LessEqual()
+    cast = P.Cast()
+    dtype = P.DType()
+    shape = P.Shape()
+    select = P.Select()
+    checktype = inner.IsSubClass()
+
+    if not checktype(dtype(input_x), mstype.float_):
+        input_x = cast(input_x, mstype.float32)
+    nan = F.fill(dtype(input_x), shape(input_x), np.nan)
+    neginf = F.fill(dtype(input_x), shape(input_x), -np.inf)
+    neg_x = less(input_x, 0.0)
+    nonpos_x = lessequal(input_x, 0.0)
+    log_x = log(input_x)
+    result = select(
+        nonpos_x, neginf, log_x)
+    return select(neg_x, nan, result)
+
+
+def log_generic_with_check(x):
+    """
+    log generic with input check
+    """
+    _check_tensor(x, "the input of log_generic")
+    return log_generic(x)
+
+
+def log1p_generic(x):
+    """
+    Log1p ops on GPU device or when device_target == GPU.
+    """
+    return log_generic(x + 1.0)
+
+
+def broadcast_to(x, target):
+    """
+    Broadcast x to the shape of target.
+    """
+    shape = P.Shape()
+    if shape(x) == shape(target):
+        return x
+    return P.BroadcastTo(shape(target))(x)