98 lines
3.0 KiB
Python
98 lines
3.0 KiB
Python
import logging
|
|
import onnx
|
|
import onnx.helper
|
|
|
|
from . import helper
|
|
from .pass_shape_inference import infer_shapes
|
|
|
|
logger = logging.getLogger("kneronnxopt.shape_optimization")
|
|
|
|
|
|
def broadcast_shape_for_Expand_checker(g, node):
|
|
# Check if the input shape and output shape have the same number of dimensions.
|
|
input_shape = helper.get_shape_from_value_info(
|
|
helper.find_value_info_everywhere_by_name(g, node.input[0])
|
|
)
|
|
output_shape = helper.get_shape_from_value_info(
|
|
helper.find_value_info_everywhere_by_name(g, node.output[0])
|
|
)
|
|
if len(input_shape) != len(output_shape):
|
|
return True
|
|
return False
|
|
|
|
|
|
def broadcast_shape_for_Expand(g, node, idx):
|
|
"""
|
|
Broadcast shape for Expand node.
|
|
|
|
:param node: the Expand node
|
|
:return: the broadcasted shape
|
|
"""
|
|
# Broadcast the input shape to the same number of dimensions as the output shape.
|
|
input_shape = helper.get_shape_from_value_info(
|
|
helper.find_value_info_everywhere_by_name(g, node.input[0])
|
|
)
|
|
output_shape = helper.get_shape_from_value_info(
|
|
helper.find_value_info_everywhere_by_name(g, node.output[0])
|
|
)
|
|
new_input_shape = [1] * (len(output_shape) - len(input_shape)) + input_shape
|
|
# Create reshape node to broadcast the input shape.
|
|
constant_node = onnx.helper.make_node(
|
|
"Constant",
|
|
[],
|
|
[node.name + "_shape"],
|
|
name=node.name + "_shape",
|
|
value=onnx.helper.make_tensor(
|
|
node.name + "_shape",
|
|
onnx.TensorProto.INT64,
|
|
dims=[len(new_input_shape)],
|
|
vals=new_input_shape,
|
|
),
|
|
)
|
|
reshape_node = onnx.helper.make_node(
|
|
"Reshape",
|
|
[node.input[0], node.name + "_shape"],
|
|
[node.input[0] + "_reshape"],
|
|
name=node.name + "_reshape",
|
|
)
|
|
g.node.insert(idx, constant_node)
|
|
g.node.insert(idx + 1, reshape_node)
|
|
helper.replace_node_input(node, node.input[0], node.input[0] + "_reshape")
|
|
return 3
|
|
|
|
|
|
pattern_registry = {
|
|
"Expand": [(broadcast_shape_for_Expand_checker, broadcast_shape_for_Expand)],
|
|
}
|
|
|
|
|
|
def run_shape_optimization_passes(model):
|
|
# Traverse the graph
|
|
i = 0
|
|
while i < len(model.graph.node):
|
|
helper.sys_still_alive_animate_print(
|
|
"[Working on Shape Opt ({status}%)]... ".format(
|
|
status=str(int(i * 100.0 / len(model.graph.node)))
|
|
)
|
|
)
|
|
|
|
node = model.graph.node[i]
|
|
if node.op_type in pattern_registry:
|
|
handled = False
|
|
for checker, handler in pattern_registry[node.op_type]:
|
|
# Check and handle the pattern
|
|
if checker(model.graph, node):
|
|
logger.debug(
|
|
f"[Doing Shape Opt] {handler.__name__} (op:{node.name})"
|
|
)
|
|
offset = handler(model.graph, node, i)
|
|
i += offset
|
|
handled = True
|
|
break
|
|
if handled:
|
|
continue
|
|
i += 1
|
|
|
|
model = infer_shapes(model)
|
|
return model
|