saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-22 19:23:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Useless Cast removal after converting model from float32 to float16 (#12871)

Browse source
This commit is contained in:
Tianlei Wu 2022-09-12 11:07:33 -07:00 committed by GitHub
parent d8636c2be8
commit 30ebc9e00a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 66 additions and 110 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

70
onnxruntime/python/tools/transformers/fusion_utils.py
View file

@ -105,45 +105,6 @@ class FusionUtils:
else:
return value == expected_value
def get_dtype(self, shape_infer_helper, input_or_output_name: str) -> int:
"""Get data type of an input or output.
Args:
shape_infer_helper (SymbolicShapeInferenceHelper): object of symbolic shape inference
input_or_output_name (str): name of input or output
Returns:
int: tensor data type
"""
dtype = self.model.get_dtype(input_or_output_name)
if dtype is not None:
return dtype
if shape_infer_helper:
tensor_proto = shape_infer_helper.known_vi_[input_or_output_name]
if tensor_proto.type.tensor_type.HasField("elem_type"):
return tensor_proto.type.tensor_type.elem_type
return None
def remove_cascaded_cast_nodes(self):
"""Remove Cast node that are overrided by another Cast node like --> Cast --> Cast -->
Note that this shall be used carefully since it might introduce semantic change.
For example, float -> int -> float could get different value than the original float value.
So, it is recommended to used only in post-processing of mixed precision conversion.
"""
removed_count = 0
for node in self.model.nodes():
if node.op_type == "Cast":
parent = self.model.get_parent(node, 0)
if parent and parent.op_type == "Cast":
node.input[0] = parent.input[0]
removed_count += 1
if removed_count > 0:
logger.info(f"Removed {removed_count} cascaded Cast nodes")
self.model.prune_graph()
def remove_identity_nodes(self):
"""Remove Identity nodes, except those right before graph output."""
nodes_to_remove = []
@ -157,34 +118,11 @@ class FusionUtils:
self.model.remove_nodes(nodes_to_remove)
logger.info(f"Removed {len(nodes_to_remove)} Identity nodes")
def remove_cascaded_cast_nodes(self):
self.model.remove_cascaded_cast_nodes()
def remove_useless_cast_nodes(self):
"""Remove cast nodes that are not needed: input and output has same data type."""
shape_infer = self.model.infer_runtime_shape(update=True)
if shape_infer is None:
return
nodes_to_remove = []
for node in self.model.nodes():
if node.op_type == "Cast":
input_dtype = self.get_dtype(shape_infer, node.input[0])
output_dtype = self.get_dtype(shape_infer, node.output[0])
if input_dtype and input_dtype == output_dtype:
nodes_to_remove.append(node)
if nodes_to_remove:
graph_input_names = set(self.model.get_graphs_input_names())
graph_output_names = set(self.model.get_graphs_output_names())
for node in nodes_to_remove:
if bool(set(node.output) & graph_output_names):
if not bool(set(node.input) & graph_input_names):
self.model.replace_output_of_all_nodes(node.input[0], node.output[0])
else:
continue
else:
self.model.replace_input_of_all_nodes(node.output[0], node.input[0])
self.model.remove_node(node)
logger.info(f"Removed {len(nodes_to_remove)} Cast nodes with output type same as input")
self.model.remove_useless_cast_nodes()
def remove_useless_reshape_nodes(self):
"""Remove reshape node that is not needed based on symbolic shape inference: input and output has same shape"""

5
onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py
View file

@ -24,7 +24,6 @@ sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
from benchmark_helper import Precision
from float16 import float_to_float16_max_diff
from fusion_utils import FusionUtils
from io_binding_helper import IOBindingHelper
from onnx_model import OnnxModel
from torch_onnx_export_helper import torch_onnx_export
@ -599,10 +598,6 @@ class Gpt2Helper:
logger.info(f"auto_mixed_precision parameters: {parameters}")
onnx_model.convert_float_to_float16(use_symbolic_shape_infer=True, **parameters)
fusion_utils = FusionUtils(onnx_model)
fusion_utils.remove_cascaded_cast_nodes()
fusion_utils.remove_useless_cast_nodes()
return parameters
@staticmethod

5
onnxruntime/python/tools/transformers/models/t5/t5_helper.py
View file

@ -20,7 +20,6 @@ from onnxruntime import InferenceSession
sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
from float16 import float_to_float16_max_diff
from fusion_utils import FusionUtils
from onnx_model import OnnxModel
from optimizer import optimize_model
@ -218,10 +217,6 @@ class T5Helper:
logger.info(f"auto_mixed_precision parameters: {parameters}")
onnx_model.convert_float_to_float16(use_symbolic_shape_infer=True, **parameters)
fusion_utils = FusionUtils(onnx_model)
fusion_utils.remove_cascaded_cast_nodes()
fusion_utils.remove_useless_cast_nodes()
return parameters
@staticmethod

96
onnxruntime/python/tools/transformers/onnx_model.py
View file

@ -8,7 +8,7 @@ import os
import sys
from collections import deque
from pathlib import Path
from typing import Dict, List, Tuple
from typing import Dict, List, Optional, Tuple
import numpy as np
from onnx import AttributeProto, GraphProto, ModelProto, NodeProto, TensorProto, helper, numpy_helper, save_model
@ -24,9 +24,9 @@ class OnnxModel:
def initialize(self, model):
self.model: ModelProto = model
self._node_name_suffix: Dict[str, int] = {} # key is node name prefix, value is the last suffix generated
self.shape_infer_helper = None
self.enable_shape_infer = True
self.all_graphs = None
self.shape_infer_helper: SymbolicShapeInferenceHelper = None
self.enable_shape_infer: bool = True
self.all_graphs: Optional[List[GraphProto]] = None
def disable_shape_inference(self):
self.enable_shape_infer = False
@ -503,6 +503,62 @@ class OnnxModel:
return value
return None
def remove_cascaded_cast_nodes(self):
"""Remove Cast node that are followed by another Cast node like --> Cast --> Cast -->
Note that this shall be used carefully since it might introduce semantic change.
For example, float -> int -> float could get different value than the original float value.
So, it is recommended to used only in post-processing of mixed precision conversion.
"""
removed_count = 0
for node in self.nodes():
if node.op_type == "Cast":
parent = self.get_parent(node, 0)
if parent and parent.op_type == "Cast":
node.input[0] = parent.input[0]
removed_count += 1
if removed_count > 0:
logger.info("Removed %d cascaded Cast nodes", removed_count)
self.prune_graph()
def remove_useless_cast_nodes(self):
"""Remove cast nodes that are not needed: input and output has same data type."""
shape_infer = self.infer_runtime_shape(update=True)
if shape_infer is None:
logger.info(f"Skip removing useless cast nodes since shape inference failed.")
return
def get_data_type(input_or_output_name):
dtype = self.get_dtype(input_or_output_name)
if dtype:
return dtype
if shape_infer.known_vi_[input_or_output_name].type.tensor_type.HasField("elem_type"):
return shape_infer.known_vi_[input_or_output_name].type.tensor_type.elem_type
return None
nodes_to_remove = []
for node in self.nodes():
if node.op_type == "Cast":
input_dtype = get_data_type(node.input[0])
output_dtype = get_data_type(node.output[0])
if input_dtype and input_dtype == output_dtype:
nodes_to_remove.append(node)
if nodes_to_remove:
graph_input_names = set(self.get_graphs_input_names())
graph_output_names = set(self.get_graphs_output_names())
for node in nodes_to_remove:
if bool(set(node.output) & graph_output_names):
if not bool(set(node.input) & graph_input_names):
self.replace_output_of_all_nodes(node.input[0], node.output[0])
else:
continue
else:
self.replace_input_of_all_nodes(node.output[0], node.input[0])
self.remove_node(node)
logger.info("Removed %d Cast nodes with output type same as input", len(nodes_to_remove))
def convert_model_float32_to_float16(self, cast_input_output=True):
logger.warning(
"The function convert_model_float32_to_float16 is deprecated. Use convert_float_to_float16 instead!"
@ -562,37 +618,9 @@ class OnnxModel:
fp16_model = convert_float_to_float16(model, **parameters)
self.initialize(fp16_model)
# Convert_float_to_float16 might add Cast(to=10) --> Cast(to=1) when two consequent nodes are computed in FP32.
# Below are post-processing that removes those Cast nodes.
# Remove first Cast nodes in path like --> Cast --> Cast -->
nodes_to_remove = []
for node in self.nodes():
if node.op_type == "Cast":
parent = self.get_parent(node, 0)
if parent and parent.op_type == "Cast":
if self.get_children(parent) == 1: # cannot be removed if its output is used by multiple nodes
self.replace_input_of_all_nodes(parent.output[0], parent.input[0])
nodes_to_remove.append(parent)
self.remove_cascaded_cast_nodes()
# Remove the second cast node.
for node in self.nodes():
if (
node.op_type == "Cast"
and OnnxModel.get_node_attribute(node, "to") == int(TensorProto.FLOAT)
and self.get_dtype(node.input[0]) == int(TensorProto.FLOAT)
):
if self.find_graph_output(node.output[0]):
self.replace_output_of_all_nodes(node.input[0], node.output[0])
else:
self.replace_input_of_all_nodes(node.output[0], node.input[0])
nodes_to_remove.append(node)
self.remove_nodes(nodes_to_remove)
if nodes_to_remove:
self.prune_graph()
print(f"removed {len(nodes_to_remove)} Cast nodes from float16 model")
self.remove_useless_cast_nodes()
def create_node_name(self, op_type, name_prefix=None):
"""Create a unique node name that starts with a prefix (default is operator type).