From f33759883c20f5b616e6af2346dcf3a2b0679b6a Mon Sep 17 00:00:00 2001
From: Yufeng Li <liyufeng1987@gmail.com>
Date: Wed, 29 Jan 2020 18:21:43 -0800
Subject: [PATCH] keep output of bert model as float32 when converting to
 float16 (#2913)

---
 onnxruntime/python/tools/bert/OnnxModel.py | 43 +++++++++++++++++-----
 1 file changed, 34 insertions(+), 9 deletions(-)

diff --git a/onnxruntime/python/tools/bert/OnnxModel.py b/onnxruntime/python/tools/bert/OnnxModel.py
index 685587b6f0..dfb8f6be10 100644
--- a/onnxruntime/python/tools/bert/OnnxModel.py
+++ b/onnxruntime/python/tools/bert/OnnxModel.py
@@ -70,6 +70,17 @@ class OnnxModel:
         for node in self.model.graph.node:
             OnnxModel.replace_node_input(node, old_input_name, new_input_name)
 
+    @staticmethod
+    def replace_node_output(node, old_output_name, new_output_name):
+        assert isinstance(old_output_name, str) and isinstance(new_output_name, str)
+        for j in range(len(node.output)):
+            if node.output[j] == old_output_name:
+                node.output[j] = new_output_name
+
+    def replace_output_of_all_nodes(self, old_output_name, new_output_name):
+        for node in self.model.graph.node:
+            OnnxModel.replace_node_output(node, old_output_name, new_output_name)
+
     def get_initializer(self,name):
         for tensor in self.model.graph.initializer:
             if tensor.name == name:
@@ -246,14 +257,6 @@ class OnnxModel:
         graph = self.model.graph
         initializers = graph.initializer
 
-        for input_value_info in graph.input:
-            if input_value_info.type.tensor_type.elem_type == 1:
-                input_value_info.type.tensor_type.elem_type = 10
-
-        for output_value_info in graph.output:
-            if output_value_info.type.tensor_type.elem_type == 1:
-                output_value_info.type.tensor_type.elem_type = 10
-
         for initializer in initializers:
             if initializer.data_type == 1:
                 initializer.CopyFrom(numpy_helper.from_array(numpy_helper.to_array(initializer).astype(np.float16), initializer.name))
@@ -267,7 +270,29 @@ class OnnxModel:
                 for att in node.attribute:
                     if att.name == 'to' and att.i == 1:
                         att.CopyFrom(onnx.helper.make_attribute("to", 10))
-   
+
+        for input_value_info in graph.input:
+            if input_value_info.type.tensor_type.elem_type == TensorProto.FLOAT:
+                initializer = self.get_initializer(input_value_info.name)
+                if initializer is not None: # for compatibility for old converter/exporter
+                    input_value_info.type.tensor_type.elem_type = 10
+                else:
+                    cast_input = input_value_info.name
+                    cast_output = input_value_info.name + '_float16'
+                    self.replace_input_of_all_nodes(cast_input, cast_output)
+                    cast_node = onnx.helper.make_node('Cast', inputs=[cast_input], outputs=[cast_output])
+                    cast_node.attribute.extend([onnx.helper.make_attribute("to", int(TensorProto.FLOAT16))])
+                    self.add_node(cast_node)
+
+        for output_value_info in graph.output:
+            if output_value_info.type.tensor_type.elem_type == TensorProto.FLOAT:
+                cast_input = output_value_info.name + '_float16'
+                cast_output = output_value_info.name
+                self.replace_output_of_all_nodes(cast_output, cast_input)
+                cast_node = onnx.helper.make_node('Cast', inputs=[cast_input], outputs=[cast_output])
+                cast_node.attribute.extend([onnx.helper.make_attribute("to", int(TensorProto.FLOAT))])
+                self.add_node(cast_node)
+
     # create a new name for node
     def create_node_name(self, op_type, name_prefix=None):
         if op_type in self.node_name_counter: