You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
370 lines
11 KiB
370 lines
11 KiB
# 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.
|
|
# ============================================================================
|
|
"""
|
|
Test nn.Distribution.
|
|
|
|
Including Normal Distribution and Bernoulli Distribution.
|
|
"""
|
|
import pytest
|
|
import numpy as np
|
|
|
|
import mindspore.nn as nn
|
|
from mindspore import dtype
|
|
from mindspore import Tensor
|
|
|
|
def test_normal_shape_errpr():
|
|
"""
|
|
Invalid shapes.
|
|
"""
|
|
with pytest.raises(ValueError):
|
|
nn.Normal([[2.], [1.]], [[2.], [3.], [4.]], dtype=dtype.float32)
|
|
|
|
def test_no_arguments():
|
|
"""
|
|
No args passed in during initialization.
|
|
"""
|
|
n = nn.Normal()
|
|
assert isinstance(n, nn.Distribution)
|
|
b = nn.Bernoulli()
|
|
assert isinstance(b, nn.Distribution)
|
|
|
|
def test_with_arguments():
|
|
"""
|
|
Args passed in during initialization.
|
|
"""
|
|
n = nn.Normal([3.0], [4.0], dtype=dtype.float32)
|
|
assert isinstance(n, nn.Distribution)
|
|
b = nn.Bernoulli([0.3, 0.5], dtype=dtype.int32)
|
|
assert isinstance(b, nn.Distribution)
|
|
|
|
class NormalProb(nn.Cell):
|
|
"""
|
|
Normal distribution: initialize with mean/sd.
|
|
"""
|
|
def __init__(self):
|
|
super(NormalProb, self).__init__()
|
|
self.normal = nn.Normal(3.0, 4.0, dtype=dtype.float32)
|
|
|
|
def construct(self, value):
|
|
x = self.normal('prob', value)
|
|
y = self.normal('log_prob', value)
|
|
return x, y
|
|
|
|
def test_normal_prob():
|
|
"""
|
|
Test pdf/log_pdf: passing value through construct.
|
|
"""
|
|
net = NormalProb()
|
|
value = Tensor([0.5, 1.0], dtype=dtype.float32)
|
|
pdf, log_pdf = net(value)
|
|
assert isinstance(pdf, Tensor)
|
|
assert isinstance(log_pdf, Tensor)
|
|
|
|
class NormalProb1(nn.Cell):
|
|
"""
|
|
Normal distribution: initialize without mean/sd.
|
|
"""
|
|
def __init__(self):
|
|
super(NormalProb1, self).__init__()
|
|
self.normal = nn.Normal()
|
|
|
|
def construct(self, value, mean, sd):
|
|
x = self.normal('prob', value, mean, sd)
|
|
y = self.normal('log_prob', value, mean, sd)
|
|
return x, y
|
|
|
|
def test_normal_prob1():
|
|
"""
|
|
Test pdf/logpdf: passing mean/sd, value through construct.
|
|
"""
|
|
net = NormalProb1()
|
|
value = Tensor([0.5, 1.0], dtype=dtype.float32)
|
|
mean = Tensor([0.0], dtype=dtype.float32)
|
|
sd = Tensor([1.0], dtype=dtype.float32)
|
|
pdf, log_pdf = net(value, mean, sd)
|
|
assert isinstance(pdf, Tensor)
|
|
assert isinstance(log_pdf, Tensor)
|
|
|
|
class NormalProb2(nn.Cell):
|
|
"""
|
|
Normal distribution: initialize with mean/sd.
|
|
"""
|
|
def __init__(self):
|
|
super(NormalProb2, self).__init__()
|
|
self.normal = nn.Normal(3.0, 4.0, dtype=dtype.float32)
|
|
|
|
def construct(self, value, mean, sd):
|
|
x = self.normal('prob', value, mean, sd)
|
|
y = self.normal('log_prob', value, mean, sd)
|
|
return x, y
|
|
|
|
def test_normal_prob2():
|
|
"""
|
|
Test pdf/log_pdf: passing mean/sd through construct.
|
|
Overwrite original mean/sd.
|
|
"""
|
|
net = NormalProb2()
|
|
value = Tensor([0.5, 1.0], dtype=dtype.float32)
|
|
mean = Tensor([0.0], dtype=dtype.float32)
|
|
sd = Tensor([1.0], dtype=dtype.float32)
|
|
pdf, log_pdf = net(value, mean, sd)
|
|
assert isinstance(pdf, Tensor)
|
|
assert isinstance(log_pdf, Tensor)
|
|
|
|
class BernoulliProb(nn.Cell):
|
|
"""
|
|
Bernoulli distribution: initialize with probs.
|
|
"""
|
|
def __init__(self):
|
|
super(BernoulliProb, self).__init__()
|
|
self.bernoulli = nn.Bernoulli(0.5, dtype=dtype.int32)
|
|
|
|
def construct(self, value):
|
|
return self.bernoulli('prob', value)
|
|
|
|
class BernoulliLogProb(nn.Cell):
|
|
"""
|
|
Bernoulli distribution: initialize with probs.
|
|
"""
|
|
def __init__(self):
|
|
super(BernoulliLogProb, self).__init__()
|
|
self.bernoulli = nn.Bernoulli(0.5, dtype=dtype.int32)
|
|
|
|
def construct(self, value):
|
|
return self.bernoulli('log_prob', value)
|
|
|
|
|
|
def test_bernoulli_prob():
|
|
"""
|
|
Test pmf/log_pmf: passing value through construct.
|
|
"""
|
|
net = BernoulliProb()
|
|
value = Tensor([1, 0, 1, 0, 1], dtype=dtype.float32)
|
|
pmf = net(value)
|
|
assert isinstance(pmf, Tensor)
|
|
|
|
def test_bernoulli_log_prob():
|
|
"""
|
|
Test pmf/log_pmf: passing value through construct.
|
|
"""
|
|
net = BernoulliLogProb()
|
|
value = Tensor([1, 0, 1, 0, 1], dtype=dtype.float32)
|
|
log_pmf = net(value)
|
|
assert isinstance(log_pmf, Tensor)
|
|
|
|
class BernoulliProb1(nn.Cell):
|
|
"""
|
|
Bernoulli distribution: initialize without probs.
|
|
"""
|
|
def __init__(self):
|
|
super(BernoulliProb1, self).__init__()
|
|
self.bernoulli = nn.Bernoulli()
|
|
|
|
def construct(self, value, probs):
|
|
return self.bernoulli('prob', value, probs)
|
|
|
|
class BernoulliLogProb1(nn.Cell):
|
|
"""
|
|
Bernoulli distribution: initialize without probs.
|
|
"""
|
|
def __init__(self):
|
|
super(BernoulliLogProb1, self).__init__()
|
|
self.bernoulli = nn.Bernoulli()
|
|
|
|
def construct(self, value, probs):
|
|
return self.bernoulli('log_prob', value, probs)
|
|
|
|
|
|
def test_bernoulli_prob1():
|
|
"""
|
|
Test pmf/log_pmf: passing probs through construct.
|
|
"""
|
|
net = BernoulliProb1()
|
|
value = Tensor([1, 0, 1, 0, 1], dtype=dtype.float32)
|
|
probs = Tensor([0.3], dtype=dtype.float32)
|
|
pmf = net(value, probs)
|
|
assert isinstance(pmf, Tensor)
|
|
|
|
def test_bernoulli_log_prob1():
|
|
"""
|
|
Test pmf/log_pmf: passing probs through construct.
|
|
"""
|
|
net = BernoulliLogProb1()
|
|
value = Tensor([1, 0, 1, 0, 1], dtype=dtype.float32)
|
|
probs = Tensor([0.3], dtype=dtype.float32)
|
|
log_pmf = net(value, probs)
|
|
assert isinstance(log_pmf, Tensor)
|
|
|
|
class BernoulliProb2(nn.Cell):
|
|
"""
|
|
Bernoulli distribution: initialize with probs.
|
|
"""
|
|
def __init__(self):
|
|
super(BernoulliProb2, self).__init__()
|
|
self.bernoulli = nn.Bernoulli(0.5)
|
|
|
|
def construct(self, value, probs):
|
|
return self.bernoulli('prob', value, probs)
|
|
|
|
class BernoulliLogProb2(nn.Cell):
|
|
"""
|
|
Bernoulli distribution: initialize with probs.
|
|
"""
|
|
def __init__(self):
|
|
super(BernoulliLogProb2, self).__init__()
|
|
self.bernoulli = nn.Bernoulli(0.5)
|
|
|
|
def construct(self, value, probs):
|
|
return self.bernoulli('log_prob', value, probs)
|
|
|
|
|
|
def test_bernoulli_prob2():
|
|
"""
|
|
Test pmf/log_pmf: passing probs/value through construct.
|
|
Overwrite original probs.
|
|
"""
|
|
net = BernoulliProb2()
|
|
value = Tensor([1, 0, 1, 0, 1], dtype=dtype.float32)
|
|
probs = Tensor([0.3], dtype=dtype.float32)
|
|
pmf = net(value, probs)
|
|
assert isinstance(pmf, Tensor)
|
|
|
|
def test_bernoulli_log_prob2():
|
|
"""
|
|
Test pmf/log_pmf: passing probs/value through construct.
|
|
Overwrite original probs.
|
|
"""
|
|
net = BernoulliLogProb2()
|
|
value = Tensor([1, 0, 1, 0, 1], dtype=dtype.float32)
|
|
probs = Tensor([0.3], dtype=dtype.float32)
|
|
log_pmf = net(value, probs)
|
|
assert isinstance(log_pmf, Tensor)
|
|
|
|
|
|
class NormalKl(nn.Cell):
|
|
"""
|
|
Test class: kl_loss of Normal distribution.
|
|
"""
|
|
def __init__(self):
|
|
super(NormalKl, self).__init__()
|
|
self.n = nn.Normal(np.array([3.0]), np.array([4.0]), dtype=dtype.float32)
|
|
|
|
def construct(self, x_, y_):
|
|
return self.n('kl_loss', 'Normal', x_, y_)
|
|
|
|
class BernoulliKl(nn.Cell):
|
|
"""
|
|
Test class: kl_loss between Bernoulli distributions.
|
|
"""
|
|
def __init__(self):
|
|
super(BernoulliKl, self).__init__()
|
|
self.b = nn.Bernoulli(0.7, dtype=dtype.int32)
|
|
|
|
def construct(self, x_):
|
|
return self.b('kl_loss', 'Bernoulli', x_)
|
|
|
|
def test_kl():
|
|
"""
|
|
Test kl_loss function.
|
|
"""
|
|
nor_net = NormalKl()
|
|
mean_b = np.array([1.0]).astype(np.float32)
|
|
sd_b = np.array([1.0]).astype(np.float32)
|
|
mean = Tensor(mean_b, dtype=dtype.float32)
|
|
sd = Tensor(sd_b, dtype=dtype.float32)
|
|
loss = nor_net(mean, sd)
|
|
assert isinstance(loss, Tensor)
|
|
|
|
ber_net = BernoulliKl()
|
|
probs_b = Tensor([0.3], dtype=dtype.float32)
|
|
loss = ber_net(probs_b)
|
|
assert isinstance(loss, Tensor)
|
|
|
|
|
|
class NormalKlNoArgs(nn.Cell):
|
|
"""
|
|
Test class: kl_loss of Normal distribution.
|
|
No args during initialization.
|
|
"""
|
|
def __init__(self):
|
|
super(NormalKlNoArgs, self).__init__()
|
|
self.n = nn.Normal(dtype=dtype.float32)
|
|
|
|
def construct(self, x_, y_, w_, v_):
|
|
return self.n('kl_loss', 'Normal', x_, y_, w_, v_)
|
|
|
|
class BernoulliKlNoArgs(nn.Cell):
|
|
"""
|
|
Test class: kl_loss between Bernoulli distributions.
|
|
No args during initialization.
|
|
"""
|
|
def __init__(self):
|
|
super(BernoulliKlNoArgs, self).__init__()
|
|
self.b = nn.Bernoulli(dtype=dtype.int32)
|
|
|
|
def construct(self, x_, y_):
|
|
return self.b('kl_loss', 'Bernoulli', x_, y_)
|
|
|
|
def test_kl_no_args():
|
|
"""
|
|
Test kl_loss function.
|
|
"""
|
|
nor_net = NormalKlNoArgs()
|
|
mean_b = np.array([1.0]).astype(np.float32)
|
|
sd_b = np.array([1.0]).astype(np.float32)
|
|
mean_a = np.array([2.0]).astype(np.float32)
|
|
sd_a = np.array([3.0]).astype(np.float32)
|
|
mean_b = Tensor(mean_b, dtype=dtype.float32)
|
|
sd_b = Tensor(sd_b, dtype=dtype.float32)
|
|
mean_a = Tensor(mean_a, dtype=dtype.float32)
|
|
sd_a = Tensor(sd_a, dtype=dtype.float32)
|
|
loss = nor_net(mean_b, sd_b, mean_a, sd_a)
|
|
assert isinstance(loss, Tensor)
|
|
|
|
ber_net = BernoulliKlNoArgs()
|
|
probs_b = Tensor([0.3], dtype=dtype.float32)
|
|
probs_a = Tensor([0.7], dtype=dtype.float32)
|
|
loss = ber_net(probs_b, probs_a)
|
|
assert isinstance(loss, Tensor)
|
|
|
|
|
|
|
|
class NormalBernoulli(nn.Cell):
|
|
"""
|
|
Test class: basic mean/sd function.
|
|
"""
|
|
def __init__(self):
|
|
super(NormalBernoulli, self).__init__()
|
|
self.n = nn.Normal(3.0, 4.0, dtype=dtype.float32)
|
|
self.b = nn.Bernoulli(0.5, dtype=dtype.int32)
|
|
|
|
def construct(self):
|
|
normal_mean = self.n('mean')
|
|
normal_sd = self.n('sd')
|
|
bernoulli_mean = self.b('mean')
|
|
bernoulli_sd = self.b('sd')
|
|
return normal_mean, normal_sd, bernoulli_mean, bernoulli_sd
|
|
|
|
def test_bascis():
|
|
"""
|
|
Test mean/sd functionality of Normal and Bernoulli.
|
|
"""
|
|
net = NormalBernoulli()
|
|
normal_mean, normal_sd, bernoulli_mean, bernoulli_sd = net()
|
|
assert isinstance(normal_mean, Tensor)
|
|
assert isinstance(normal_sd, Tensor)
|
|
assert isinstance(bernoulli_mean, Tensor)
|
|
assert isinstance(bernoulli_sd, Tensor)
|