saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-27 20:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

keep output of bert model as float32 when converting to float16 (#2913)

Browse source
This commit is contained in:
Yufeng Li 2020-01-29 18:21:43 -08:00 committed by GitHub
parent a319ad29b3
commit f33759883c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 34 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
onnxruntime/python/tools/bert/OnnxModel.py
View file

@ -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: