From 2b134c72e6ffdd38d408efcf30bb49111bdc030f Mon Sep 17 00:00:00 2001 From: Kittipat Virochsiri Date: Tue, 24 Jul 2018 11:47:00 -0700 Subject: [PATCH] Add interface to provide blob types to shape&type inference (#9643) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/9643 Current map interface assumes float data type, which is not always correct. Reviewed By: kennyhorror Differential Revision: D8455784 fbshipit-source-id: b94a31267760f7f97c15aa4b03008affc347fd10 --- caffe2/core/operator.cc | 25 ++++++++++++++++ caffe2/core/operator.h | 5 ++++ caffe2/operators/atomic_ops.cc | 3 +- caffe2/operators/map_ops.cc | 3 +- caffe2/operators/segment_reduction_op.h | 7 +++++ .../operator_test/shape_inference_test.py | 30 +++++++++++++++++++ caffe2/python/pybind_state.cc | 27 +++++++++++++++++ caffe2/python/workspace.py | 10 +++++-- 8 files changed, 106 insertions(+), 4 deletions(-) diff --git a/caffe2/core/operator.cc b/caffe2/core/operator.cc index 169e730125a..2d3b8994bda 100644 --- a/caffe2/core/operator.cc +++ b/caffe2/core/operator.cc @@ -571,6 +571,31 @@ TensorShapes InferBlobShapesAndTypesFromMap( return InferBlobShapesAndTypes(blob_desc, nets); } +TensorShapes InferBlobShapesAndTypesFromMap( + const CaffeMap>& blob_dimensions, + const CaffeMap& blob_types, + const vector& nets) { + CaffeMap blob_desc; + // Populate shapes from known blobs + for (const auto& blob : blob_dimensions) { + TensorShape tp; + for (auto d : blob.second) { + CAFFE_ENFORCE_GE(d, 0, blob.first); + tp.add_dims(d); + } + auto blob_type = blob_types.find(blob.first); + if (blob_type == blob_types.end()) { + LOG(WARNING) << "Missing type of " << blob.first + << "; assuming to be UNDEFINED"; + tp.set_data_type(TensorProto_DataType_UNDEFINED); + } else { + tp.set_data_type(blob_type->second); + } + blob_desc[blob.first] = tp; + } + return InferBlobShapesAndTypes(blob_desc, nets); +} + std::map> ValidateTensorDevices( OperatorBase& op, const OperatorDef& op_def) { diff --git a/caffe2/core/operator.h b/caffe2/core/operator.h index 26bb02415d3..325ccd3761a 100644 --- a/caffe2/core/operator.h +++ b/caffe2/core/operator.h @@ -968,6 +968,11 @@ TensorShapes InferBlobShapesAndTypesFromMap( const CaffeMap>& blob_dimensions, const vector& nets); +TensorShapes InferBlobShapesAndTypesFromMap( + const CaffeMap>& blob_dimensions, + const CaffeMap& blob_types, + const vector& nets); + std::map> ValidateTensorDevices( OperatorBase& op, const OperatorDef& op_def); diff --git a/caffe2/operators/atomic_ops.cc b/caffe2/operators/atomic_ops.cc index 111302446d9..31a4dd659f7 100644 --- a/caffe2/operators/atomic_ops.cc +++ b/caffe2/operators/atomic_ops.cc @@ -93,7 +93,8 @@ OPERATOR_SCHEMA(CreateMutex) .NumInputs(0) .NumOutputs(1) .SetDoc("Creates an unlocked mutex and returns it in a unique_ptr blob.") - .Output(0, "mutex_ptr", "Blob containing a std::unique_ptr."); + .Output(0, "mutex_ptr", "Blob containing a std::unique_ptr.") + .ScalarType(TensorProto_DataType_UNDEFINED); OPERATOR_SCHEMA(AtomicFetchAdd) .NumInputs(3) diff --git a/caffe2/operators/map_ops.cc b/caffe2/operators/map_ops.cc index 4862165bea1..f2da863d15a 100644 --- a/caffe2/operators/map_ops.cc +++ b/caffe2/operators/map_ops.cc @@ -58,7 +58,8 @@ OPERATOR_SCHEMA(CreateMap) .SetDoc("Create an empty map blob") .Arg("key_dtype", "Key's TensorProto::DataType (default INT32)") .Arg("value_dtype", "Value's TensorProto::DataType (default INT32)") - .Output(0, "map blob", "Blob reference to the map"); + .Output(0, "map blob", "Blob reference to the map") + .ScalarType(TensorProto_DataType_UNDEFINED); OPERATOR_SCHEMA(KeyValueToMap) .NumInputs(2) diff --git a/caffe2/operators/segment_reduction_op.h b/caffe2/operators/segment_reduction_op.h index 24303165637..1d51692ac71 100644 --- a/caffe2/operators/segment_reduction_op.h +++ b/caffe2/operators/segment_reduction_op.h @@ -2007,6 +2007,13 @@ i.e. `len(LENGTHS)`. Other dimensions are inherited from the input tensor. "OUTPUT", "Aggregated output tensor. Has the first dimension of K " "(the number of segments)."); + schema.TensorInferenceFunction( + [](const OperatorDef&, const std::vector& input_types) { + std::vector out(1); + out[0] = input_types[0]; + out[0].set_dims(0, input_types[Reducer::kInputCount + 1].dims(0)); + return out; + }); ReducerDef::PopulateSchema(schema); } using Reducer = typename ReducerDef::template Reducer; diff --git a/caffe2/python/operator_test/shape_inference_test.py b/caffe2/python/operator_test/shape_inference_test.py index 59eb899c97c..d08231d09a3 100644 --- a/caffe2/python/operator_test/shape_inference_test.py +++ b/caffe2/python/operator_test/shape_inference_test.py @@ -431,6 +431,36 @@ class TestShapeInference(test_util.TestCase): self.assertEqual(shapes['E'], [10, 23, 9, 10]) self.assertEqual(shapes['G'], [10, 23, 9, 2, 10]) + def testConcatInt32(self): + net = core.Net("concat") + + net.Concat(["A", "B"], ["C", "splits"], axis=1) + net.Concat(["C", "D"], ["E"], order="NCHW") + net.Concat(["E", "F"], ["G"], add_axis=1, order="NHWC") + (shapes, types) = workspace.InferShapesAndTypes( + [net], + blob_dimensions={ + 'A': [10, 12, 9, 10], + 'B': [10, 9, 9, 10], + 'D': [10, 2, 9, 10], + 'F': [10, 23, 9, 10] + }, + blob_types={ + 'A': core.DataType.INT32, + 'B': core.DataType.INT32, + 'D': core.DataType.INT32, + 'F': core.DataType.INT32, + } + ) + self.assertEqual(shapes['C'], [10, 21, 9, 10]) + self.assertEqual(shapes['splits'], [2]) + self.assertEqual(shapes['E'], [10, 23, 9, 10]) + self.assertEqual(shapes['G'], [10, 23, 9, 2, 10]) + self.assertEqual(types['C'], core.DataType.INT32) + self.assertEqual(types['splits'], core.DataType.INT32) + self.assertEqual(types['E'], core.DataType.INT32) + self.assertEqual(types['G'], core.DataType.INT32) + def testSqueeze(self): net = core.Net("sq") net.Squeeze(["data"], ["data_squeezed"], dims=[3, 1]) diff --git a/caffe2/python/pybind_state.cc b/caffe2/python/pybind_state.cc index 48b316c771d..70bc635193f 100644 --- a/caffe2/python/pybind_state.cc +++ b/caffe2/python/pybind_state.cc @@ -1354,6 +1354,33 @@ void addGlobalMethods(py::module& m) { auto blob_info = InferBlobShapesAndTypesFromMap(blob_dimensions, nets_ptr); + std::string protob; + CAFFE_ENFORCE(blob_info.SerializeToString(&protob)); + return py::bytes(protob); + }); + m.def( + "infer_shapes_and_types_from_map", + [](const std::vector& net_protos, + const std::map> blob_dimensions, + const std::map int_blob_types) { + // Parse protobuffers to NetDefs + std::vector> nets; + std::vector nets_ptr; + for (auto proto : net_protos) { + std::unique_ptr def(new NetDef()); + CAFFE_ENFORCE(def->ParseFromString(proto)); + nets_ptr.push_back(def.get()); + nets.push_back(std::move(def)); + } + std::map blob_types; + for (auto blob_type : int_blob_types) { + blob_types[blob_type.first] = + static_cast(blob_type.second); + } + + auto blob_info = InferBlobShapesAndTypesFromMap( + blob_dimensions, blob_types, nets_ptr); + std::string protob; CAFFE_ENFORCE(blob_info.SerializeToString(&protob)); return py::bytes(protob); diff --git a/caffe2/python/workspace.py b/caffe2/python/workspace.py index 8c3c1729db2..c033e0684bb 100644 --- a/caffe2/python/workspace.py +++ b/caffe2/python/workspace.py @@ -236,7 +236,8 @@ def RunPlanInBackground(plan_or_step): return C.run_plan_in_background(StringifyProto(plan_or_step)) -def InferShapesAndTypes(nets, blob_dimensions=None, nets_proto=False): +def InferShapesAndTypes(nets, blob_dimensions=None, nets_proto=False, + blob_types=None): """Infers the shapes and types for the specified nets. Inputs: @@ -253,11 +254,16 @@ def InferShapesAndTypes(nets, blob_dimensions=None, nets_proto=False): else: net_protos = [StringifyProto(n.Proto()) for n in nets] if blob_dimensions is None: + assert blob_types is None blobdesc_prototxt = C.infer_shapes_and_types_from_workspace(net_protos) - else: + elif blob_types is None: blobdesc_prototxt = C.infer_shapes_and_types_from_map( net_protos, blob_dimensions ) + else: + blobdesc_prototxt = C.infer_shapes_and_types_from_map( + net_protos, blob_dimensions, blob_types + ) blobdesc_proto = caffe2_pb2.TensorShapes() blobdesc_proto.ParseFromString(blobdesc_prototxt) shapes = {}