Infer shape from data in Constant nodes (#1667)

shinh · houseroad · commit 0ab3c951bde3 · 2018-12-04T00:10:33.000-08:00
diff --git a/onnx/shape_inference/implementation.cc b/onnx/shape_inference/implementation.cc
@@ -106,9 +106,22 @@ static void InferShapesImpl(
       valueTypesByName[vi.name()] = vi.mutable_type();
   }
 
-  std::unordered_map<std::string, const TensorProto*> initializersByName;
+  std::unordered_map<std::string, const TensorProto*> inputDataByName;
   for (const auto& tp : g->initializer()) {
-    initializersByName[tp.name()] = &tp;
+    inputDataByName[tp.name()] = &tp;
+  }
+  // Collect data from constant nodes.
+  for (const auto& n : g->node()) {
+      if (n.op_type() != "Constant" || n.output().size() != 1) {
+          continue;
+      }
+      for (const auto& attr : n.attribute()) {
+          if (attr.name() == "value" &&
+              attr.type() == AttributeProto::TENSOR &&
+              attr.has_t()) {
+              inputDataByName[n.output(0)] = &attr.t();
+          }
+      }
   }
 
   for (auto& n : *g->mutable_node()) {
@@ -122,7 +135,7 @@ static void InferShapesImpl(
     const auto schema =
         schema_registry->GetSchema(n.op_type(), domain_version, n.domain());
     InferenceContextImpl ctx(
-        n, valueTypesByName, initializersByName, &graphInferenceContext);
+        n, valueTypesByName, inputDataByName, &graphInferenceContext);
     if (!schema) {
       if (nullptr == func_registry) {
         continue;
diff --git a/onnx/test/shape_inference_test.py b/onnx/test/shape_inference_test.py
@@ -222,6 +222,17 @@ def test_reshape_static_shape_inferred(self):  # type: () -> None
             initializer=[make_tensor('shape', TensorProto.INT64, (3,), (0, 3, -1))])
         self._assert_inferred(graph, [make_tensor_value_info('y', TensorProto.UINT8, (2, 3, 4))])
 
+    def test_reshape_static_shape_constant(self):  # type: () -> None
+        graph = self._make_graph(
+            [('x', TensorProto.UINT8, (2, 4, 3))],
+            [make_node("Constant", [], ['shape'],
+                       value=make_tensor('shape', TensorProto.INT64, (2,), (3, 8))),
+             make_node("Reshape", ['x', 'shape'], ['y'])],
+            [])
+        self._assert_inferred(graph, [
+            make_tensor_value_info('shape', TensorProto.INT64, (2,)),
+            make_tensor_value_info('y', TensorProto.UINT8, (3, 8))])
+
     def test_upsample(self):  # type: () -> None
         graph = self._make_graph(
             [('x', TensorProto.INT32, (2, 4, 3, 5)),
