Impl Shape op for mkldnn (#15266)

Summary:
Impl Shape op for mkldnn
Pull Request resolved: https://github.com/pytorch/pytorch/pull/15266

Differential Revision: D13804558

Pulled By: yinghai

fbshipit-source-id: 8a35f608c23973d7a15c3d645aee4059eb55f245

diff --git a/caffe2/ideep/operators/shape_op.cc b/caffe2/ideep/operators/shape_op.cc
new file mode 100644 (file)
index 0000000..cdf3b36
--- /dev/null
+++ b/caffe2/ideep/operators/shape_op.cc
@@ -0,0 +1,68 @@
+#include <caffe2/ideep/ideep_utils.h>
+
+namespace caffe2 {
+
+// RecordShapeOp records the shape of the input tensor to a vector of int. You
+// mostly don't need this operator explicitly, and it is mostly used in the
+// autodiff process.
+class IDEEPShapeOp : public IDEEPOperator {
+ public:
+  USE_IDEEP_DEF_ALIASES();
+  USE_IDEEP_OPERATOR_FUNCTIONS();
+
+  IDEEPShapeOp(const OperatorDef& operator_def, Workspace* ws)
+      : IDEEPOperator(operator_def, ws),
+        axes_(OperatorBase ::GetRepeatedArgument<int>("axes")) {}
+
+  bool RunOnDevice() override {
+    int numDims = 0;
+    int numAxes = axes_.size();
+    vector<int64_t> dims;
+    const char* data_dims = nullptr;
+    auto* output = OperatorBase::Output<Tensor>(OUTPUT, CPU);
+
+    if (OperatorBase::InputBlob(DATA).template IsType<itensor>()) {
+      auto& data = Input(DATA);
+      numDims = data.ndims();
+      auto idims = data.get_dims();
+      dims.assign(idims.begin(), idims.end());
+      data_dims = reinterpret_cast<const char*>(dims.data());
+    } else {
+      auto& data = OperatorBase::Input<Tensor>(DATA, CPU);
+      numDims = data.dim();
+      data_dims = reinterpret_cast<const char*>(data.sizes().data());
+    }
+
+    if (numAxes == 0) {
+      output->Resize(numDims);
+      int64_t* output_data = output->template mutable_data<int64_t>();
+      context_.CopyBytesSameDevice(
+          numDims * sizeof(int64_t), data_dims, output_data);
+      return true;
+    }
+
+    output->Resize(numAxes);
+    auto out = reinterpret_cast<char*>(output->template mutable_data<int64_t>());
+    for (int i = 0; i < numAxes; i++) {
+      auto axis = axes_[i];
+      CAFFE_ENFORCE_LT(axis, numDims, "Axis out of range");
+      CAFFE_ENFORCE_GE(axis, 0, "Each axis should be non-negative");
+      context_.CopyBytesSameDevice(
+          sizeof(int64_t), data_dims + axis * sizeof(int64_t), out);
+      out += sizeof(int64_t);
+    }
+
+    return true;
+  }
+
+ private:
+  vector<int> axes_;
+
+  INPUT_TAGS(DATA);
+  OUTPUT_TAGS(OUTPUT);
+};
+
+
+REGISTER_IDEEP_OPERATOR(Shape, IDEEPShapeOp);
+
+} // namespace caffe2

diff --git a/caffe2/python/ideep/shape_op_test.py b/caffe2/python/ideep/shape_op_test.py
new file mode 100644 (file)
index 0000000..d87ae54
--- /dev/null
+++ b/caffe2/python/ideep/shape_op_test.py
@@ -0,0 +1,89 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import unittest
+import hypothesis.strategies as st
+from hypothesis import given, settings
+import numpy as np
+from caffe2.proto import caffe2_pb2
+from caffe2.python import core, workspace
+from caffe2.python.transformations import optimizeForIDEEP
+import caffe2.python.hypothesis_test_util as hu
+import caffe2.python.ideep_test_util as mu
+
+
+@unittest.skipIf(not workspace.C.use_mkldnn, "No MKLDNN support.")
+class ShapeTest(hu.HypothesisTestCase):
+    @given(n=st.integers(1, 128),
+           c=st.integers(1, 128),
+           h=st.integers(1, 128),
+           w=st.integers(1, 128),
+           **mu.gcs)
+    def test_shape(self, n, c, h, w, gc, dc):
+        op0 = core.CreateOperator(
+            "Shape",
+            ["X0"],
+            ["Y0"],
+            device_option=dc[0]
+        )
+        op1 = core.CreateOperator(
+            "Shape",
+            ["X1"],
+            ["Y1"],
+            device_option=dc[1]
+        )
+        X = np.random.rand(n, c, h, w).astype(np.float32) - 0.5
+        workspace.FeedBlob('X0', X, dc[0])
+        workspace.FeedBlob('X1', X, dc[1])
+        workspace.RunOperatorOnce(op0)
+        workspace.RunOperatorOnce(op1)
+        Y0 = workspace.FetchBlob('Y0')
+        Y1 = workspace.FetchBlob('Y1')
+
+        if not np.allclose(Y0, Y1, atol=0, rtol=0):
+            print(Y1.flatten())
+            print(Y0.flatten())
+            print(np.max(np.abs(Y1 - Y0)))
+            self.assertTrue(False)
+
+    @given(n=st.integers(1, 128),
+           c=st.integers(1, 128),
+           h=st.integers(1, 128),
+           w=st.integers(1, 128),
+           axes=st.lists(st.integers(0, 3), min_size=1, max_size=3),
+           **mu.gcs)
+    def test_shape_with_axes(self, n, c, h, w, axes, gc, dc):
+        axes = list(set(axes)).sort()
+        op0 = core.CreateOperator(
+            "Shape",
+            ["X0"],
+            ["Y0"],
+            axes = axes,
+            device_option=dc[0]
+        )
+        op1 = core.CreateOperator(
+            "Shape",
+            ["X1"],
+            ["Y1"],
+            axes = axes,
+            device_option=dc[1]
+        )
+        X = np.random.rand(n, c, h, w).astype(np.float32) - 0.5
+        workspace.FeedBlob('X0', X, dc[0])
+        workspace.FeedBlob('X1', X, dc[1])
+        workspace.RunOperatorOnce(op0)
+        workspace.RunOperatorOnce(op1)
+        Y0 = workspace.FetchBlob('Y0')
+        Y1 = workspace.FetchBlob('Y1')
+
+        if not np.allclose(Y0, Y1, atol=0, rtol=0):
+            print(Y1.flatten())
+            print(Y0.flatten())
+            print(np.max(np.abs(Y1 - Y0)))
+            self.assertTrue(False)
+
+
+if __name__ == "__main__":
+    unittest.main()