pytorch

1

2

3

4

5

6
import unittest
7
import hypothesis.strategies as st
8
from hypothesis import given, settings
9
import numpy as np
10
from caffe2.python import core, workspace
11
import caffe2.python.hypothesis_test_util as hu
12
import caffe2.python.mkl_test_util as mu
13

14

15
@unittest.skipIf(not workspace.C.has_mkldnn,
16
                 "Skipping as we do not have mkldnn.")
17
class MKLConvTest(hu.HypothesisTestCase):
18
    @given(stride=st.integers(1, 3),
19
           pad=st.integers(0, 3),
20
           kernel=st.integers(3, 5),
21
           size=st.integers(8, 8),
22
           input_channels=st.integers(1, 3),
23
           output_channels=st.integers(1, 3),
24
           batch_size=st.integers(1, 3),
25
           **mu.gcs)
26
    @settings(max_examples=2, deadline=100)
27
    def test_mkl_convolution(self, stride, pad, kernel, size,
28
                             input_channels, output_channels,
29
                             batch_size, gc, dc):
30
        op = core.CreateOperator(
31
            "Conv",
32
            ["X", "w", "b"],
33
            ["Y"],
34
            stride=stride,
35
            pad=pad,
36
            kernel=kernel,
37
        )
38
        X = np.random.rand(
39
            batch_size, input_channels, size, size).astype(np.float32) - 0.5
40
        w = np.random.rand(
41
                output_channels, input_channels, kernel, kernel) \
42
            .astype(np.float32) - 0.5
43
        b = np.random.rand(output_channels).astype(np.float32) - 0.5
44

45
        inputs = [X, w, b]
46
        self.assertDeviceChecks(dc, op, inputs, [0])
47

48

49
if __name__ == "__main__":
50
    import unittest
51
    unittest.main()
52