Merge pull request #88 from pengwa/fix_output_dtype

fix the graph output node dtype issue

Author: guschmue

tf2onnx/graph.py

@@ -300,6 +300,10 @@ def get_dtype(self, name):
         """Get dtype for node."""
         return self._dtypes.get(name)
 
+    def set_dtype(self, name, val):
+        """Set dtype for node."""
+        self._dtypes[name] = val
+
     def get_shape(self, name):
         """Get shape for node."""
         assert isinstance(name, str)
@@ -390,13 +394,11 @@ def make_model(self, doc, input_names, output_names, optimize=True):
         # create output_tensor_values
         output_tensor_values = []
         for name in output_names:
-            op = self.get_node_by_name(name)
-            if op:
-                dtype = op.dtype
-                if not dtype:
-                    continue
-                v = helper.make_tensor_value_info(name, dtype, self.get_shape(name))
-                output_tensor_values.append(v)
+            dtype = self.get_dtype(name);
+            if not dtype:
+                raise ValueError("cannot found the output dtype for " + name)
+            v = helper.make_tensor_value_info(name, dtype, self.get_shape(name))
+            output_tensor_values.append(v)
 
         # update attributes
         ops = []

tf2onnx/tfonnx.py

@@ -898,6 +898,9 @@ def topk_op(ctx, node, name, args):
     node.set_attr("k", k[0])
     node.type = "TopK"
     ctx.remove_input(node, node.input[1])
+
+    # the second of TopK operator must be INT64 per ONNX requires. 
+    ctx.set_dtype(name + ":1", onnx_pb.TensorProto.INT64)
     return node