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
|
import unittest
import hypothesis.strategies as st
from hypothesis import given
import numpy as np
from caffe2.python import core, workspace
import caffe2.python.hypothesis_test_util as hu
import caffe2.python.mkl_test_util as mu
@unittest.skipIf(
not workspace.C.has_mkldnn, "Skipping as we do not have mkldnn."
)
class MKLConcatTest(hu.HypothesisTestCase):
@given(
batch_size=st.integers(1, 10),
channel_splits=st.lists(st.integers(1, 10), min_size=1, max_size=3),
height=st.integers(1, 10),
width=st.integers(1, 10),
**mu.gcs
)
def test_mkl_concat(
self, batch_size, channel_splits, height, width, gc, dc
):
Xs = [
np.random.rand(batch_size, channel,
height, width).astype(np.float32)
for channel in channel_splits
]
op = core.CreateOperator(
"Concat",
["X_{}".format(i) for i in range(len(Xs))],
["concat_result", "split_info"],
order="NCHW",
)
self.assertDeviceChecks(dc, op, Xs, [0])
if __name__ == "__main__":
import unittest
unittest.main()
|
