pytorch

1

2

3

4

5

6
from caffe2.python import core, workspace
7
from hypothesis import assume, given, settings
8
import caffe2.python.hypothesis_test_util as hu
9
import caffe2.python.serialized_test.serialized_test_util as serial
10
import hypothesis.strategies as st
11
import numpy as np
12

13

14
class TestReductionOps(serial.SerializedTestCase):
15

16
    @serial.given(n=st.integers(5, 8), **hu.gcs)
17
    def test_elementwise_sum(self, n, gc, dc):
18
        X = np.random.rand(n).astype(np.float32)
19

20
        def sum_op(X):
21
            return [np.sum(X)]
22

23
        op = core.CreateOperator(
24
            "SumElements",
25
            ["X"],
26
            ["y"]
27
        )
28

29
        self.assertReferenceChecks(
30
            device_option=gc,
31
            op=op,
32
            inputs=[X],
33
            reference=sum_op,
34
        )
35

36
        self.assertGradientChecks(
37
            device_option=gc,
38
            op=op,
39
            inputs=[X],
40
            outputs_to_check=0,
41
            outputs_with_grads=[0],
42
        )
43

44
    @given(n=st.integers(5, 8), **hu.gcs)
45
    @settings(deadline=10000)
46
    def test_elementwise_int_sum(self, n, gc, dc):
47
        X = np.random.rand(n).astype(np.int32)
48

49
        def sum_op(X):
50
            return [np.sum(X)]
51

52
        op = core.CreateOperator(
53
            "SumElementsInt",
54
            ["X"],
55
            ["y"]
56
        )
57

58
        self.assertReferenceChecks(
59
            device_option=gc,
60
            op=op,
61
            inputs=[X],
62
            reference=sum_op,
63
        )
64

65
    @given(n=st.integers(1, 65536),
66
           dtype=st.sampled_from([np.float32, np.float16]),
67
           **hu.gcs)
68
    @settings(deadline=10000)
69
    def test_elementwise_sqrsum(self, n, dtype, gc, dc):
70
        if dtype == np.float16:
71
            # fp16 is only supported with CUDA/HIP
72
            assume(gc.device_type == workspace.GpuDeviceType)
73
            dc = [d for d in dc if d.device_type == workspace.GpuDeviceType]
74

75
        X = np.random.rand(n).astype(dtype)
76

77
        def sumsqr_op(X):
78
            return [np.sum(X * X)]
79

80
        op = core.CreateOperator(
81
            "SumSqrElements",
82
            ["X"],
83
            ["y"]
84
        )
85

86
        threshold = 0.01 if dtype == np.float16 else 0.005
87

88
        self.assertReferenceChecks(
89
            device_option=gc,
90
            op=op,
91
            inputs=[X],
92
            reference=sumsqr_op,
93
            threshold=threshold,
94
        )
95

96
    @given(n=st.integers(5, 8), **hu.gcs)
97
    def test_elementwise_avg(self, n, gc, dc):
98
        X = np.random.rand(n).astype(np.float32)
99

100
        def avg_op(X):
101
            return [np.mean(X)]
102

103
        op = core.CreateOperator(
104
            "SumElements",
105
            ["X"],
106
            ["y"],
107
            average=1
108
        )
109

110
        self.assertReferenceChecks(
111
            device_option=gc,
112
            op=op,
113
            inputs=[X],
114
            reference=avg_op,
115
        )
116

117
        self.assertGradientChecks(
118
            device_option=gc,
119
            op=op,
120
            inputs=[X],
121
            outputs_to_check=0,
122
            outputs_with_grads=[0],
123
        )
124

125
    @serial.given(batch_size=st.integers(1, 3),
126
           m=st.integers(1, 3),
127
           n=st.integers(1, 4),
128
           **hu.gcs)
129
    def test_rowwise_max(self, batch_size, m, n, gc, dc):
130
        X = np.random.rand(batch_size, m, n).astype(np.float32)
131

132
        def rowwise_max(X):
133
            return [np.max(X, axis=2)]
134

135
        op = core.CreateOperator(
136
            "RowwiseMax",
137
            ["x"],
138
            ["y"]
139
        )
140

141
        self.assertReferenceChecks(
142
            device_option=gc,
143
            op=op,
144
            inputs=[X],
145
            reference=rowwise_max,
146
        )
147

148
    @serial.given(batch_size=st.integers(1, 3),
149
           m=st.integers(1, 3),
150
           n=st.integers(1, 4),
151
           **hu.gcs)
152
    def test_columnwise_max(self, batch_size, m, n, gc, dc):
153
        X = np.random.rand(batch_size, m, n).astype(np.float32)
154

155
        def columnwise_max(X):
156
            return [np.max(X, axis=1)]
157

158
        op = core.CreateOperator(
159
            "ColwiseMax",
160
            ["x"],
161
            ["y"]
162
        )
163

164
        self.assertReferenceChecks(
165
            device_option=gc,
166
            op=op,
167
            inputs=[X],
168
            reference=columnwise_max,
169
        )
170

171
        # Test shape inference logic
172
        net = core.Net("test_shape_inference")
173
        workspace.FeedBlob("x", X)
174
        output = net.ColwiseMax(["x"], ["y"])
175
        (shapes, types) = workspace.InferShapesAndTypes([net])
176
        workspace.RunNetOnce(net)
177

178
        self.assertEqual(shapes[output], list(workspace.blobs[output].shape))
179
        self.assertEqual(shapes[output], [X.shape[0]] + [X.shape[2]])
180
        self.assertEqual(types[output], core.DataType.FLOAT)
181