/
test_walker_alias.py
51 lines (40 loc) · 1.45 KB
/
test_walker_alias.py
1
2
3
4
5
6
7
8
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
import unittest
import numpy
from chainer import backend
from chainer.backends import cuda
from chainer import testing
from chainer.testing import attr
from chainer import utils
class TestWalkerAlias(unittest.TestCase):
def setUp(self):
self.ps = numpy.array([5, 3, 4, 1, 2], dtype=numpy.int32)
self.sampler = utils.WalkerAlias(self.ps)
def check_sample(self):
counts = numpy.zeros(len(self.ps), numpy.float32)
for _ in range(1000):
vs = self.sampler.sample((4, 3))
numpy.add.at(counts, cuda.to_cpu(vs), 1)
counts /= (1000 * 12)
counts *= sum(self.ps)
testing.assert_allclose(self.ps, counts, atol=0.1, rtol=0.1)
def test_sample_cpu(self):
assert not self.sampler.use_gpu
self.check_sample()
@attr.gpu
def test_sample_gpu(self):
with testing.assert_warns(DeprecationWarning):
self.sampler.to_gpu()
assert self.sampler.use_gpu
assert isinstance(self.sampler._device.device, cuda.Device)
self.check_sample()
@attr.gpu
def test_to_cpu(self):
with testing.assert_warns(DeprecationWarning):
self.sampler.to_gpu()
with testing.assert_warns(DeprecationWarning):
self.sampler.to_cpu()
assert not self.sampler.use_gpu
assert isinstance(
self.sampler._device, backend.CpuDevice)
self.check_sample()
testing.run_module(__name__, __file__)