MindSpore概率采样库系列（一）：手把手教你开发概率分布

2020/11/06

MindSpore概率采样库系列（一）：手把手教你开发概率分布

MindSpore概率采样库系列共有两篇

1.手把手教你开发概率分布

2.手把手教你开发概率映射

概率分布

为了描述生活中的各种随机事件，数学家提出了随机变量的概念。随机变量的值是不确定的，但是他们取不同值的可能性满足一定的概率分布。我们可以用一些函数描述一个随机变量的不同输出值的可能性，比如概率密度函数描述这个随机变量的输出值在某个确定的取值点附近的可能性，以及累计分布函数描述这个随机变量的输出值不大于某个确定的取值点的可能性。另外，我们可以用一些统计量去描述一个随机变量的某些高度概括的概率性质，比如期望(mean)是随机变量所有可能输出值的概率乘以其结果的总和，以及众数（mode）指一组数据中出现次数最多的数据值。

概率分布是一个随机变量的概率性质的数学抽象，包括我们可以对随机变量所有可以做的操作和计算，比如给定一个输出值计算对应的概率，又如从随机分布中取一个样本。概率分布是深度概率编程的基础。例如在前文中提到的变分自编码器，当从隐向量空间解码时，我们需要对一个高斯分布进行中随机采样，从而生成一个向量作为解码器的输入。而在最后做优化的时候，我们要计算两个概率分布的KL散度作为优化的目标之一。

Mindspore 中为用户提供了概率分布采用库，其中包括了多样的基础概率分布及其基本函数计算，也为 Mindspore 的贝叶斯概率网络建模能力提供了基础。今天向大家介绍 Mindspore 里的概率分布基类 Distribution, 并以 Laplace 分布为例，教大家如何给MindSpore增加一个概率分布。

前期准备

为了开发一个概率分布，我们需要做一些前期的准备。以本文介绍的 Laplace 分布为例，开发前我们需要了解他的一些数学性质，包括：

参数形式。一个概率分布可能是由多种参数形式的。例如对于一个指数分布，我们可以使用 rate 作为参数，也可以用 rate 的倒数 scale 作为参数来描述这个分布。这里我们选择 Laplace 分布的 loc 和 scale 作为参数描述这个分布的性质。
分布函数。在确定了参数表达之后，我们需要了解在如此参数设定下，概率分布的相关分布函数的表达式，例如概率密度函数，累计分布函数，生存函数等等。以 Laplace 分布为例，当我们确定了他的参数为 loc 和 scale 后，我们可以确定他的概率密度函数的表达式为 f(x) = 1/(2 * scale) exp(-|x - loc| / scale) 。
统计量的计算。对于一些常用的统计量，例如期望和方差，我们需要了解在我们参数设定下如何去计算他们。以 Laplace 分布为例，当我们确定了他的参数为 loc 和 scale 后， Laplace 分布的期望为 loc，而方差为 2*scale^2 .
抽样方法。如何从一个概率分布中抽样一直是一个重要的研究课题。高效的抽样能让有效的帮助模型的训练和推理。对于连续性分布，我们可以采用如下抽样方法：假设F(x)是某个分布的累计分布函数，那么如果u是一个来自均匀分布Uniform(0, 1)的抽样，那么x = F^{-1}(u)就是该分布的一个抽样。下面我们将采用这个抽样方法。

构建函数 (Constructor)

在概率分布层中， Distribution 是所有概率分布的基类。通常来说，有两种类型的分布。第一种类型的概率分布，例如 Normal 和 Uniform，他们直接继承 Distribution，初始化参数包括概率分布的特征参数 dist_spec_args，概率分布的类型 dtype 以及概率分布的名称 name。另一种概率分布继承自TransformedDistribution。 TransformedDistribution 继承自 Distribution ，初始化参数包括一个 Bijector 和一个 Distribution，例如通过 Exp bijector 和 Normal distribution 实现的 LogNormal 分布。今天介绍的 Laplace 属于第一种简单类型的概率分布。下面我们开始吧！

首先我们新建一个继承自 Distribution 的名为 Laplace的派生类，定义 __init__ 函数。

class Laplace(Distribution):
    def __init__(self,
                loc=0.,
                scale=1.,
                dtype=mstype.float32,
                name='Laplace'):

这里我们的参数包括 loc, scale, 是 Laplace 分布的特征参数 dist_spec_args，默认值为分别为0和1。分布类型dtype，因为Laplace 是连续分布，我们设置默认值为float32。分布名称 name 默认值为 'Laplace'。

dist_spec_args 详解

文中提到的dist_spec_args是某一分布的特征参数。例如，mean 和 std 可作为 Normal 分布的特征参数，而 rate 可作为 Expnential 分布的特征参数。这里，loc, scale 为 Laplace 的特征参数，他们可以是具体的数值，也可以是空值（None）。你可能有一个问题，既然 dist_spec_args 是完整定义一个分布的必要值，为什么其值可以空呢？

当一个概率分布的初始化特征某个参数为空时，此实例可理解为一个dummy distribution。在概率编程中，很多时候，分布的特征参数跟其他函数有关或为变量，此时，我们希望仅初始化分布的类型。因此我们支持dist_spec_args在初始化时为 None, 在调用具体函数时传入。需强调的是，为空的 dist_spec_args 必须在调用具体的函数时以 args 或者 kwargs的方式传入，否则函数将报错。具体参考后面文章中对 _check_param_type 函数的讨论。

有了Laplace的类型定义，我们开始处理传入的参数吧！

class Laplace(Distribution):
    def __init__(self, loc=None， scale=None, dtype=mstype.float32, name='Laplace'):
        param = dict(locals())
        param['param_dict']={"loc": loc, "scale", scale}
        valid_type = mstype.float_type
        Validator.check_type(type(self.__name__), dtype, valid, type)
        super(Laplace, self).__init__(seed, dtype, name, param)

在调用基类构造函数前，即Distribution类构造函数前，我们需要：

利用 Python build_in locals 将函数初始化时用到的参数存入 param 字典中。
在param 中加入以 'param_dict' 为key, 值为 {“dist_spec_arg1_name” : dist_spec_arg1, ...}的元素
用 Validator 检查 dtype的正确性

完成以上步骤后就我们可以调用 Distribution 基类的构造函数啦！

Distribution 基类构造函数

Distribution 继承Cell，构造函数参数包括:

seed: 随机抽样种子。
dtype: 分布的数据类型。
name: 分布的名称。
param: 分布初始化时所用参数。

在Distribution基类构造函数中，完成了检查name, seed 参数的合法性，计算并设置该分布的属性。

Distribution 构造函数中初始化的Property 和公有属性

Distribution 类的公有的Property包括name, dtype, seed, parameters, is_scalar_bacth, batch_shape, broadcast_shape。这些属性都在基类构造函数中被初始化。。

self.parameter_type

self.parameter_type的检查由 mindspore.nn.probabilsity.distribution._utils.utils 中set_param_type实现。 set_param_type 接收一个参数字典，和一个 hint_type。一般来说，hint_type设为self.dtype，即分布的类型。我们希望self.parameter_type为浮点数类型。所以对于离散函数来说，hint_type默认为 mindspore.float32。当一个分布有多个参数，且以_np.array_ 或 _Tensor_的形式传入时，参数的dtype必须一致，否则将会报错。出于统一的后端支持的原因，当传入的参数 dtype 是某些后端运算不支持的类型，例如 float64 时, 参数将会被自动类型转换为 hint_type。

is_scalar_batch, batch_shape, broadcast_shape

若初始化时，param['param_dict']中值皆为 int/float/None 是， is_scalar_batch 为 True。对于非TransformedDistribution 的分布来说，batch_shpae 等于 broadcast_shape，是所有参数广播后的 shape。若参数不能广播，将会报错。

私有函数指针

除了一些Property 和公有属性外，基类构造函数中还初始化了_call_prob等函数指针。该指针指向派生类的_prob函数或基类的wrapper functions，这一部分会在第三节中详细介绍。

为派生类添加属性

完成基类初始化后，我们需要在子类中初始化子类属性, 通常为分布的 dist_spec_args, 这一步需调用 Distribution 中基类函数 _add_parameter，将传入的参数类型转换为类型为self.parameter_type的 Tensor，并将该参数及参数名称加入到 self.default_parameters 和 self.parameters_names两个 list 中。以 Laplace 为例，分布的 dist_spec_args 为 location, 和 scale, 因此，我们初始化self._loc, self._scale。 _add_parameter 完成后，self.default_parameters = [self._loc, self._scale]，self.parameters_names = ['loc', 'scale']。这两个属性在具体的概率函数实现中尤为重要，详见第三节 Log_prob 的参数。

如果你希望检查参数的正确性，mindspore.nn.probabilsity.distribution._utils.utils 中提供了很多 utility functions 可供调用。此处，我们检查 scale 的是否大于零。

self._loc = self._add_parameter(loc, 'loc')
    self._scale = self._add_parameter(scale, 'scale')
    if self._scale is not None:
        check_greater_zero(self._scale， 'scale')

另外，我们可以在构造函数外将self._loc，self._scale 设为 property，方便访问。

@property
   def loc(self):
        return self._loc
    @property
   def loc(self):
        return self._loc

最后, 在构造函数中初始化你需要用到的算子和nupmy的常数，算子通常在mindspore.ops 和mindspore.nn.probabiltiy.distribution._utils.custom_ops 下找到。custom_ops中提供更加稳定的算子。这里以四个算子为例：

self.abs = P.Abs()
    self.cast = P.Cast()
    self.exp = exp_generic
    self.log = log_generic

到这里，你已经完成了Laplace分布的构造函数，完整的代码如下：

class Laplace(Distribution):
    def __init__(self, loc=0.， scale=1., dtype=mstype.float32, name='Laplace'):
        param = dict(locals())
        param['param_dict']={"loc": loc, "scale", scale}
        valid_type = mstype.float_type
        Validator.check_type(type(self.__name__), dtype, valid, type)
        super(Laplace, self).__init__(seed, dtype, name, param)
        self._loc = self._add_parameter(loc, 'loc')
        self._scale = self._add_parameter(scale, 'scale')
        if self._scale is not None:
            check_greater_zero(self._scale， ‘scale’)
        
        self.abs = P.Abs()
        self.cast = P.Cast()
        self.const = P.ScalarToArray()
        self.exp = exp_generic
        self.log = log_generic
        self.log1p = P.Log1p()
        self.sqrt = P.Sqrt()
        self.shape = P.Shape()
        self.sign = P.Sign()
        self.squeeze = P.Squeeze(0)
        self.uniform = C.uniform
        
        self.eps = np.finfo(float).eps
    
    @property
    def loc(self):
        return self._loc
    @property
    def loc(self):
        return self._loc

实现打包函数参数的公有接口

下面我们来实现一些需要重载的公有接口， ·_get_dist_type 和 _get_dist_args。 _get_dist_type 不接收任何参数，返回一个代表函数类型的字符串。

def _get_dist_type(self):
        return "Laplace"

_get_dist_args 接收 dist_spec_args, 可打包函数以list 的形式返回, 一般形式为

def  _get_dist_type(self，dist_spec_args1=None， dist_spec_args2=None, ... )

在此函数中，需检查传入的 dist_spec_args 是否为 Tensor（可以直接调用基类已初始化的self.checktensor 算子）。如果传入 dist_spec_args 为 None, 返回初始化时的值。具体实现如下：

def _get_dist_args(self, loc=None, scale=None):
        if loc is not None:
            self.checktensor(loc, 'loc')
        else:
            loc = self.loc
        if scale is not None:
            self.checktensor(scale, 'scale')
        else:
            scale = self.scale
        return loc, scale

实现 log_prob 函数

现在我们进入到具体的函数的实现把！ Distribution 类支持的函数包括 prob, log_prob, cdf, log_cdf, survival_function, log_survival, mean, sd, var, entropy, kl_loss, cross_entropy,和 sample，由派生类函数的具体实现决定参数。调用时，基类会调用对应函数指针指向的函数，这些指针在基类构造函数中比如说，使用prob时，会调用指针self._call_prob指向的函数。 self._call_prob默认指向_prob, 所以，派生类中需重载 _prob 来支持基类 prob 计算。

完整的函数表：

prob : 概率密度函数（pdf）/ 概率质量函数（pmf）
log_prob ：对数似然函数。
cdf : 累积分布函数（cdf）。
log_cdf : 对数累积分布函数（cdf）。
survival_function : 生存函数。与 cdf 互补。
log_survival : 对数生存函数。
mean : 均值。
mode : 众数。
sd : 标准差。
var : 方差。
entropy : 熵。
kl_loss : Kullback-Leibler散度。
cross_entropy : 两个概率分布的交叉熵。
sample : 随机取样。
get_dist_type: 分布类型
get_dist_args: 分布特征参数

Wrapper Functions

然而并不是每一个函数都需要在派生类中实现， Distribution 实现了多样的 wrapper functions，完成函数间的相互转换。例如， _calc_prob_from_log_prob，通过求 log_prob 来获得 prob 的值。Distribution中该函数的实现如下：

def _calc_prob_from_log_prob(self, value, *args, **kwargs):
        return self.exp_base(self._log_prob(value, *args, **kwargs))

在 Distribution 构造函数中，_set_prob, _set_log_prob, _set_cdf 等函数通过判断派生类的函数实现情况，设定基类函数指针。举个例子，若派生类中实现了 _prob 函数，但未实现 _log_prob函数，默认地，_set_log_prob 会将 self._call_log_prob 函数指针指向 _calc_prob_fromlog__prob。因为prob, log_prob 之间可相互转换，所以派升类只需实现其中的一个函数，若两个函数都没有在派生类中实现，那么这两个函数的指针指向一个报错函数 _raise_not_implemented_error。当尝试调用_log_prob或_prob是，会出现NotImplementedError。但为了函数稳定性与计算效率，亦可同时在派生类中实现_prob, _log_prob。此时。基类函数中的_set_prob 等函数将不会选择 wrapper function，而会直接指向派生类中的函数。

今天在Laplace分布中，我们选择重载_log_prob。 prob的计算将由基类调用_calc_prob_from_log_prob进行。

log_prob 的参数

一般地，prob, log_prob, cdf, log_cdf, survival_function, log_survival 有相同的参数形式，必须传入一个为 Tensor 的 value。可选参数为分布的特征参数 dist_spec_args, 默认值为None。以Laplace 为例，可选参数为loc，scale。根据以上原则，log_prob 的函数定义如下：

def _log_prob(self, value, loc=None, scale=None):

进入函数后，我们需要对value, loc,scale 进行检查。

对value进行处理

value = self._check_value(value, "value")
    value = self.cast(value, self.dtype)

_check_value 为Distribution基类函数，检查value 是否为 Tensor 类型，如果不是Tensor将报错。为保证函数稳定性，建议将所有 input value 都转换为 float_type。这里可以将 value 类型转换为 self.dtype. 离散分布可将value 类型转换为 self.parameter_type.

检查loc, scale 这里有四种情况需要考虑，1. 初始化时 loc, scale 不为空（即 self.loc, self.scale 不为空），log_prob 时未传入 loc, scale; 2. self.loc, self.scale 不为空, log_prob 传入新的 loc, scale; 3. self.loc, self.scale 为空, log_prob 未传入 loc, scale; 4. self.loc, self.scale 为空, log_prob时传入新的 loc, scale. 对这四种请况的讨论可直接调用基类函数 _check_param_type 完成对 self.loc, self.scale, loc, scale的检查与选择。

loc, scale = self._check_param_type(loc, scale)

_check_param_type

需注意的是，在调用此函数时，loc, scale 的传入顺序需与__init__中 _add_parameter的顺序一致。例如，在__init__中，加入的顺序为 self.loc, self.scale, 此时传入的顺序亦为先 loc 后 scale。在这个函数中，会检查传入的 loc, scale 正确性，并决定使用此时传入的loc, scale ，或初始化时存入在 self.default_parameters 中的 self.loc和self.scale。

若传入的loc,scale 不为 None, 且为 Tensor, 此函数会选择使用新传入的 loc, scale。若传入的loc,scale 不为 None, 但不为 Tensor, 将会报错。若传入的loc,scale 为 None, 此函数会检查初始化时的self.loc,self.scale是否为None, 如果不为 None,函数将使用初始化时候的值。若传入的loc,scale 和初始化值都为 None,函数将报错, 即初始化时若建立了一个 dummy distribution, 必须在_log_prob等函数中通过 loc, sccale 传入参数。简而言之，若某一个 _dist_spec_args1_在初始化时为 None, 则在计算函数中需传入一个为 Tensor 的 dist_spec_args1。若某一个 _dist_spec_args1_在初始化时不为 None, 可选择传入的 dist_spec_args1 。另外，新传入的 dist_spec_args 只会在此次调用中替代 self.default_parameters, 但是不会 overwrite self.default_parameters。

需注意的是，在_check_param_type 中会完成loc与scale的广播. 例: 若 loc 为 shape 为 (1,), scale 为 (2, 1), 此函数return 后 loc将被 broadcast 成一个(2,1)的 Tensor。

检查处理完参数后，就可以开始写计算公式啦！调用 constuctor 中已初始化的算子，计算_log_prob。

z = (x - loc) / scale
return -1. \* (self.abs(z) + self.log(2. \* scale))

z = (x - loc) / scale
    return -1. * (self.abs(z) + self.log(2. * scale))

完整的 _log_prob 代码如下：

def _log_prob(self, value, loc=None, scale=None):
        value = self._check_value(value, "value")
        value = self.cast(value, self.dtype)
        loc, scale = self._check_param_type(loc, scale)
        z = (x - loc) / scale
    return -1. * (self.abs(z) + self.log(2. * scale))

在完成 _log_prob 之后，相信大家能够完成类似的_cdf， _log_cdf, _survival_function, _log_survival 函数。需要注意的是，这四个函数中，择一实现即可。

KL散度

这一节中向大家介绍 _kl_loss的写法。kl_loss计算两个分布之间的KL散度, 而在讯息系统中称为相对熵, 记为 KL(a||b)。KL散度是两个概率分布 a 和 b 差别的非对称性的度量，或者可以理解他是两个概率分布之间的距离。KL散度中的两个分布的地位是不同的，典型情况下，a 表示数据的真实分布，b 表示数据的理论分布、估计的模型分布、或 b 的近似分布。在很多机器学习的模型中会使用KL散度作为目标函数的一部分作为优化目标。

当我们使用一个概率分布计算KL散度的时候，他本身是作为KL散度计算里的 a 分布。那么不同于其他函数，kl_loss的输入则是KL散度计算中的另一个概率分布。这里，我们把分布 b 以参数列表的形式，作为必选参数传入函数中。参数列表的包括两个部分，参数的类型和参数对应的 dist_spec_args 。此处我们使用 dist 以_string_ 的形式指定参数类型，loc_b 和 scale_b 组成分布 b 的 dist_spec_args。这样我们就获得了Laplace kl_loss函数定义：

def _kl_loss(self, dist, loc_b, scale_b, loc=None, scale=None):

因为只有当分布 b 同为 Laplace 分布时我们有对应KL散度的解析解公式，我们调用 mindspore.nn.probabilsity.distribution._utils.utils 中 check_distribution_name 对 dist 判定输入的分布是我们支持的类型。

check_distribution_name（dist, "Laplace"）

之后，我们检查loc_b, scale_b 是否为 Tensor, 并将其类型转换为self.parameter_type。

loc_b = self._check_value(loc_b, 'loc_b')
    loc_b = self.cast(loc_b, self.parameter_type)
    scale_b = self._check_value(scale_b, 'scale_b')
    scale_b = self.cast(scale_b, self.parameter_type)

最后检查 loc, scale。

loc, scale = self._check_param_type(loc, scale)

完成参数检查后，就可以开始写计算公式了。

loc_diff = self.abs(loc - loc_b)
    scale_log_diff = self.log(scale) - self.log(scale_b)
    return (-1 * scale_log_diff +\
            loc_diff / scale_b - 1. +\
            self.exp(-1 * loc_diff / scale + scale_log_diff))

完整的_kl_loss代码如下：

def _kl_loss(self, dist, loc_b, scale_b, loc=None, scale=None):
        check_distribution_name（dist, 'Laplace'）
        loc_b = self._check_value(loc_b, 'loc_b')
        loc_b = self.cast(loc_b, self.parameter_type)
        scale_b = self._check_value(scale_b, 'scale_b')
        scale_b = self.cast(scale_b, self.parameter_type)
        loc, scale = self._check_param_type(loc, scale)
        loc_diff = self.abs(loc - loc_b)
        scale_log_diff = self.log(scale) - self.log(scale_b)
        return (-1 * scale_log_diff +\
            loc_diff / scale_b - 1. +\
            self.exp(-1 * loc_diff / scale + scale_log_diff))

采样 Sample

最后，我们介绍采样函数 sample 的写法。此函数接受一个类型为 tuple 的shape，默认值为空值 ()。dist_spec_args 可选择传入。类似地，我们首先进行参数检查。我们调用基类算子 _check_tuple 对shape进行检查, 并检查是否传入了新的 loc, scale。

def _sample(self, shape=(), loc=None, scale=None):
        shape = self._check_tuple(shape, "shape")
        loc, scale = self._check_param_type(loc, scale)

首先计算样本的形状, original_shape 为 [shape , batch_shape]。若传入的shape 与batch_shape 均为scalar, 即 ()，我们在此处增加一维。

original_shape = shape + self.shape(loc)
    if original_shape == ():
        sample_shape = (1,)
    else:
        sample_shape = original_shape

根据上面介绍的抽样方法，这里我们利用 uniform 抽样算子，变换得到实现 Laplace 抽样，并将得到的samples 类型转换为分布类型，即self.dtype。

low = self.const(-1.0 + self.eps)
    high = self.const(1.0)
    uniform_samples = self.uniform(sample_shape, low, high, self.seed)
    samples = (loc - scale * self.sign(uniform_samples) *
            self.log1p(-1. * self.abs(uniform_samples)))
    samples = self.cast(samples, self.dtype)

最后，若original_shape 为()，降维：

if original_shape == ():
        return self.squeeze(samples)
    return samples

完整的代码如下：

def _sample(self, shape=(), loc=None, scale=None):
        shape = self._check_tuple(shape, "shape")
        loc, scale = self._check_param_type(loc, scale)
        original_shape = shape + self.shape(loc)
        if original_shape == ():
            sample_shape = (1,)
        else:
            sample_shape = original_shape
        low = self.const(-1.0 + self.eps)
        high = self.const(1.0)
        uniform_samples = self.uniform(sample_shape, low, high, self.seed)
        samples = (loc - scale * self.sign(uniform_samples) *
                self.log1p(-1. * self.abs(uniform_samples)))
        samples = self.cast(samples, self.dtype)
        if original_shape == ():
            return self.squeeze(samples)
        return samples

Congrats!

到这里教程就结束了，恭喜你完成了第一个自定义概率分布。MindSpore作为开源社区期待你加入更多的概率分布，让我们一起来扩充概率分布类吧！

作者介绍： 作者邓洵（华为多伦多异构编译器实验室实习生）目前就读于加拿大多伦多大学，参与设计与开发MindSpore概率分布采用库，是MindSpore开源社区概率编程模块的重要代码贡献者。

学习

核心框架

大模型套件

科学计算套件

领域套件

工具

模型库

生态资源

技术学习

专区

社区组织

贡献与成长

开发者活动

互动交流

活动

动态

MindSpore概率采样库系列（一）：手把手教你开发概率分布

MindSpore概率采样库系列（一）：手把手教你开发概率分布

概率分布

前期准备

构建函数 (Constructor)

dist_spec_args 详解

Distribution 基类构造函数

Distribution 构造函数中初始化的Property 和公有属性

self.parameter_type

is_scalar_batch, batch_shape, broadcast_shape

私有函数指针

为派生类添加属性

实现打包函数参数的公有接口

实现 log_prob 函数

Wrapper Functions

log_prob 的参数

_check_param_type

KL散度

采样 Sample