|
|
|
|
|
|
|
|
|
|
| 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 MKLConvTest(hu.HypothesisTestCase): |
| @given(stride=st.integers(1, 3), |
| pad=st.integers(0, 3), |
| kernel=st.integers(3, 5), |
| size=st.integers(8, 20), |
| input_channels=st.integers(1, 16), |
| output_channels=st.integers(1, 16), |
| batch_size=st.integers(1, 3), |
| use_bias=st.booleans(), |
| group=st.integers(1, 8), |
| **mu.gcs) |
| def test_mkl_convolution(self, stride, pad, kernel, size, |
| input_channels, output_channels, |
| batch_size, use_bias, group, gc, dc): |
| op = core.CreateOperator( |
| "Conv", |
| ["X", "w", "b"] if use_bias else ["X", "w"], |
| ["Y"], |
| stride=stride, |
| pad=pad, |
| kernel=kernel, |
| group=group |
| ) |
| X = np.random.rand( |
| batch_size, input_channels * group, size, size).astype(np.float32) - 0.5 |
| w = np.random.rand( |
| output_channels * group, input_channels, kernel, kernel) \ |
| .astype(np.float32) - 0.5 |
| b = np.random.rand(output_channels * group).astype(np.float32) - 0.5 |
|
|
| inputs = [X, w, b] if use_bias else [X, w] |
| self.assertDeviceChecks(dc, op, inputs, [0]) |
|
|
|
|
| if __name__ == "__main__": |
| import unittest |
| unittest.main() |
|
|
