wehub-resource-sync
29 lines
815 B
Python
Executable File
29 lines
815 B
Python
Executable File
import unittest
|
|
import numpy as np
|
|
from prml import nn
|
|
|
|
|
|
class TestDirichlet(unittest.TestCase):
|
|
|
|
def test_dirichlet(self):
|
|
np.random.seed(1234)
|
|
obs = np.random.choice(3, 100, p=[0.2, 0.3, 0.5])
|
|
obs = np.eye(3)[obs]
|
|
a = nn.Parameter(np.zeros(3))
|
|
for _ in range(100):
|
|
a.cleargrad()
|
|
mu = nn.softmax(a)
|
|
d = nn.random.Dirichlet(np.ones(3) * 10, data=mu)
|
|
x = nn.random.Categorical(mu, data=obs)
|
|
log_posterior = x.log_pdf().sum() + d.log_pdf().sum()
|
|
log_posterior.backward()
|
|
a.value += 0.01 * a.grad
|
|
|
|
count = np.sum(obs, 0) + 10
|
|
p = count / count.sum(keepdims=True)
|
|
self.assertTrue(np.allclose(p, mu.value, 1e-2, 1e-2))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|