pytorch

1

2

3

4

5

6
import numpy as np
7

8
from hypothesis import given
9
import hypothesis.strategies as st
10

11
from caffe2.python import core
12
from caffe2.python import workspace
13
import caffe2.python.hypothesis_test_util as hu
14

15

16
class TestWeightedMultiSample(hu.HypothesisTestCase):
17
    @given(
18
        num_samples=st.integers(min_value=0, max_value=128),
19
        data_len=st.integers(min_value=0, max_value=10000),
20
        **hu.gcs_cpu_only
21
    )
22
    def test_weighted_multi_sample(self, num_samples, data_len, gc, dc):
23
        weights = np.zeros((data_len))
24
        expected_indices = []
25
        if data_len > 0:
26
            weights[-1] = 1.5
27
            expected_indices = np.repeat(data_len - 1, num_samples)
28

29
        workspace.FeedBlob("weights", weights.astype(np.float32))
30

31
        op = core.CreateOperator(
32
            "WeightedMultiSampling",
33
            ["weights"],
34
            ["sample_indices"],
35
            num_samples=num_samples,
36
        )
37
        workspace.RunOperatorOnce(op)
38
        result_indices = workspace.FetchBlob("sample_indices")
39
        np.testing.assert_allclose(expected_indices, result_indices)
40
        self.assertDeviceChecks(
41
            dc,
42
            op,
43
            [weights.astype(np.float32)],
44
            [0]
45
        )
46

47
        # test shape input
48
        shape = np.zeros((num_samples))
49
        workspace.FeedBlob("shape", shape)
50
        op2 = core.CreateOperator(
51
            "WeightedMultiSampling",
52
            ["weights", "shape"],
53
            ["sample_indices_2"]
54
        )
55
        workspace.RunOperatorOnce(op2)
56
        result_indices_2 = workspace.FetchBlob("sample_indices_2")
57
        if data_len > 0:
58
            assert len(result_indices_2) == num_samples
59
            for i in range(num_samples):
60
                assert 0 <= result_indices_2[i] < data_len
61
        else:
62
            assert len(result_indices_2) == 0
63

64
        self.assertDeviceChecks(dc, op2, [weights.astype(np.float32), shape], [0])
65

66

67
if __name__ == "__main__":
68
    import unittest
69
    unittest.main()
70