wehub-resource-sync
422 lines
19 KiB
Python
422 lines
19 KiB
Python
import inspect
|
|
import re
|
|
from typing import Dict, Tuple, Union, Literal
|
|
|
|
import onnx
|
|
import torch
|
|
from google.protobuf.internal.containers import RepeatedCompositeFieldContainer
|
|
from onnx import GraphProto, ModelProto, NodeProto, ValueInfoProto
|
|
|
|
__verbose__: bool = True
|
|
"""
|
|
Whether log information of successful operations
|
|
"""
|
|
|
|
|
|
# Whether the running torch.onnx.export() exposes the `dynamo` keyword.
|
|
# Its presence is fragmented across PyTorch versions: introduced as a
|
|
# separate `torch.onnx.dynamo_export` API in 2.1, added as a kwarg on
|
|
# `torch.onnx.export` in 2.4, removed in some intermediate releases, and
|
|
# reinstated (with a True default) in 2.9. We probe the signature once at
|
|
# import time and only pass `dynamo=False` when the kwarg actually exists.
|
|
# All ONNX graph surgery in this module is written against the TorchScript
|
|
# exporter, so we want to stay on it whenever the choice is offered.
|
|
TORCHSCRIPT_EXPORT_KWARGS: Dict[str, object] = (
|
|
{'dynamo': False}
|
|
if 'dynamo' in inspect.signature(torch.onnx.export).parameters
|
|
else {}
|
|
)
|
|
|
|
|
|
def _verbose(self, *args, sep=' ', end='\n', file=None):
|
|
if __verbose__:
|
|
print(self, *args, sep=sep, end=end, file=file)
|
|
|
|
|
|
def model_override_io_shapes(
|
|
model: ModelProto,
|
|
input_shapes: Dict[str, Tuple[Union[str, int]]] = None,
|
|
output_shapes: Dict[str, Tuple[Union[str, int]]] = None,
|
|
):
|
|
"""
|
|
Override the shapes of inputs/outputs of the model graph (in-place operation).
|
|
:param model: model to perform the operation on
|
|
:param input_shapes: a dict with keys as input/output names and values as shape tuples
|
|
:param output_shapes: the same as input_shapes
|
|
"""
|
|
def _override_shapes(
|
|
shape_list_old: RepeatedCompositeFieldContainer[ValueInfoProto],
|
|
shape_dict_new: Dict[str, Tuple[Union[str, int]]]):
|
|
for value_info in shape_list_old:
|
|
if value_info.name in shape_dict_new:
|
|
name = value_info.name
|
|
dims = value_info.type.tensor_type.shape.dim
|
|
assert len(shape_dict_new[name]) == len(dims), \
|
|
f'Number of given and existing dimensions mismatch: {name}'
|
|
for i, dim in enumerate(shape_dict_new[name]):
|
|
if isinstance(dim, int):
|
|
dims[i].dim_param = ''
|
|
dims[i].dim_value = dim
|
|
else:
|
|
dims[i].dim_value = 0
|
|
dims[i].dim_param = dim
|
|
_verbose(f'| override shape of \'{name}\' with {shape_dict_new[name]}')
|
|
|
|
if input_shapes is not None:
|
|
_override_shapes(model.graph.input, input_shapes)
|
|
if output_shapes is not None:
|
|
_override_shapes(model.graph.output, output_shapes)
|
|
|
|
|
|
def model_reorder_io_list(
|
|
model: ModelProto,
|
|
input_or_output: Literal['input', 'output'],
|
|
target_name: str,
|
|
insert_after_name: str,
|
|
):
|
|
"""
|
|
Reorder the input of the model graph by moving the target input after the specified input (in-place operation).
|
|
If the given names are not found, the operation will be ignored.
|
|
:param model: model to perform the operation on
|
|
:param input_or_output: 'input' or 'output' to specify the list to reorder
|
|
:param target_name: the name of the input to be reordered
|
|
:param insert_after_name: the name of the input to be inserted after (None for the first)
|
|
"""
|
|
def _reorder_input(input_list: RepeatedCompositeFieldContainer[ValueInfoProto]):
|
|
nonlocal input_or_output
|
|
target_idx = -1
|
|
insert_after_idx = -1
|
|
for i, value_info in enumerate(input_list):
|
|
if value_info.name == target_name:
|
|
target_idx = i
|
|
if value_info.name == insert_after_name:
|
|
insert_after_idx = i
|
|
if target_idx != -1 and insert_after_idx != -1:
|
|
target = input_list.pop(target_idx)
|
|
input_list.insert(insert_after_idx + 1, target)
|
|
_verbose(f'| reorder {input_or_output}: \'{target_name}\' after \'{insert_after_name}\'')
|
|
|
|
if input_or_output == 'input':
|
|
_reorder_input(model.graph.input)
|
|
elif input_or_output == 'output':
|
|
_reorder_input(model.graph.output)
|
|
else:
|
|
raise ValueError('Argument \'input_or_output\' should be either \'input\' or \'output\'.')
|
|
|
|
|
|
def model_add_prefixes(
|
|
model: ModelProto,
|
|
initializer_prefix=None,
|
|
value_info_prefix=None,
|
|
node_prefix=None,
|
|
dim_prefix=None,
|
|
ignored_pattern=None,
|
|
):
|
|
"""
|
|
Adds prefixes to names inside the given ONNX model graph, including sub-graphs (in-place operation).
|
|
This method is a complete version of the official onnx.compose.add_prefix API, which does not consider sub-graphs.
|
|
"""
|
|
initializers = set()
|
|
value_infos = set()
|
|
|
|
def _record_initializers_and_value_infos_recursive(subgraph):
|
|
# Record names in current graph
|
|
for initializer in subgraph.initializer:
|
|
if ignored_pattern is not None and re.match(ignored_pattern, initializer.name):
|
|
continue
|
|
initializers.add(initializer.name)
|
|
for value_info in subgraph.value_info:
|
|
if ignored_pattern is not None and re.match(ignored_pattern, value_info.name):
|
|
continue
|
|
value_infos.add(value_info.name)
|
|
for node in subgraph.node:
|
|
# For 'If' and 'Loop' nodes, do recording recursively
|
|
if node.op_type == 'If':
|
|
for attr in node.attribute:
|
|
branch = onnx.helper.get_attribute_value(attr)
|
|
_record_initializers_and_value_infos_recursive(branch)
|
|
elif node.op_type == 'Loop':
|
|
for attr in node.attribute:
|
|
if attr.name == 'body':
|
|
body = onnx.helper.get_attribute_value(attr)
|
|
_record_initializers_and_value_infos_recursive(body)
|
|
|
|
def _add_prefixes_recursive(subgraph):
|
|
# Add prefixes in current graph
|
|
if initializer_prefix is not None:
|
|
for initializer in subgraph.initializer:
|
|
if ignored_pattern is not None and re.match(ignored_pattern, initializer.name):
|
|
continue
|
|
new_name = initializer_prefix + initializer.name
|
|
_verbose('| add prefix:', initializer.name, '->', new_name)
|
|
initializer.name = new_name
|
|
|
|
for value_info in subgraph.value_info:
|
|
if dim_prefix is not None:
|
|
for dim in value_info.type.tensor_type.shape.dim:
|
|
if dim.dim_param is None or dim.dim_param == '' or \
|
|
ignored_pattern is not None and re.match(ignored_pattern, dim.dim_param):
|
|
continue
|
|
new_dim_param = dim_prefix + dim.dim_param
|
|
_verbose('| add prefix:', dim.dim_param, '->', new_dim_param)
|
|
dim.dim_param = new_dim_param
|
|
|
|
if value_info_prefix is None or \
|
|
ignored_pattern is not None and re.match(ignored_pattern, value_info.name):
|
|
continue
|
|
new_name = value_info_prefix + value_info.name
|
|
_verbose('| add prefix:', value_info.name, '->', new_name)
|
|
value_info.name = new_name
|
|
|
|
if node_prefix is not None:
|
|
for node in subgraph.node:
|
|
if ignored_pattern is not None and re.match(ignored_pattern, node.name):
|
|
continue
|
|
new_name = node_prefix + node.name
|
|
_verbose('| add prefix:', node.name, '->', new_name)
|
|
node.name = new_name
|
|
|
|
for node in subgraph.node:
|
|
# For 'If' and 'Loop' nodes, add prefixes recursively
|
|
if node.op_type == 'If':
|
|
for attr in node.attribute:
|
|
branch = onnx.helper.get_attribute_value(attr)
|
|
_add_prefixes_recursive(branch)
|
|
elif node.op_type == 'Loop':
|
|
for attr in node.attribute:
|
|
if attr.name == 'body':
|
|
body = onnx.helper.get_attribute_value(attr)
|
|
_add_prefixes_recursive(body)
|
|
|
|
# For each node, rename its inputs and outputs
|
|
for io_list in [node.input, node.output]:
|
|
for i, io_value in enumerate(io_list):
|
|
if io_value in initializers and initializer_prefix is not None:
|
|
new_value = initializer_prefix + io_value
|
|
_verbose('| add prefix:', io_value, '->', new_value)
|
|
io_list[i] = new_value
|
|
if io_value in value_infos and value_info_prefix is not None:
|
|
new_value = value_info_prefix + io_value
|
|
_verbose('| add prefix:', io_value, '->', new_value)
|
|
io_list[i] = new_value
|
|
|
|
_record_initializers_and_value_infos_recursive(model.graph)
|
|
_add_prefixes_recursive(model.graph)
|
|
|
|
|
|
def graph_fold_back_to_squeeze(graph: GraphProto):
|
|
"""
|
|
Fold the substructures of 'Shape', 'Gather', 'Equal', 'If' to one single 'Squeeze' node.
|
|
This can unify the different behaviors between aten::squeeze and onnx:Squeeze.
|
|
"""
|
|
def _graph_fold_back_to_squeeze_recursive(subgraph: GraphProto):
|
|
# Do folding in sub-graphs recursively.
|
|
for node in subgraph.node:
|
|
if node.op_type == 'If':
|
|
for attr in node.attribute:
|
|
branch = onnx.helper.get_attribute_value(attr)
|
|
_graph_fold_back_to_squeeze_recursive(branch)
|
|
elif node.op_type == 'Loop':
|
|
for attr in node.attribute:
|
|
if attr.name == 'body':
|
|
body = onnx.helper.get_attribute_value(attr)
|
|
_graph_fold_back_to_squeeze_recursive(body)
|
|
|
|
# Do folding in current graph.
|
|
i_shape = 0
|
|
while i_shape < len(subgraph.node):
|
|
if subgraph.node[i_shape].op_type == 'Shape':
|
|
shape_node = subgraph.node[i_shape]
|
|
shape_out = shape_node.output[0]
|
|
i_gather = i_shape + 1
|
|
while i_gather < len(subgraph.node):
|
|
if subgraph.node[i_gather].op_type == 'Gather' and subgraph.node[i_gather].input[0] == shape_out:
|
|
gather_node = subgraph.node[i_gather]
|
|
gather_out = gather_node.output[0]
|
|
i_equal = i_gather + 1
|
|
while i_equal < len(subgraph.node):
|
|
if subgraph.node[i_equal].op_type == 'Equal' and (
|
|
subgraph.node[i_equal].input[0] == gather_out
|
|
or subgraph.node[i_equal].input[1] == gather_out):
|
|
equal_node = subgraph.node[i_equal]
|
|
equal_out = equal_node.output[0]
|
|
i_if = i_equal + 1
|
|
while i_if < len(subgraph.node):
|
|
if subgraph.node[i_if].op_type == 'If' \
|
|
and subgraph.node[i_if].input[0] == equal_out:
|
|
# Found the substructure to be folded.
|
|
if_node = subgraph.node[i_if]
|
|
# Create 'Squeeze' node.
|
|
squeeze_node = onnx.helper.make_node(
|
|
op_type='Squeeze',
|
|
inputs=[
|
|
*list(shape_node.input),
|
|
# For ONNX opset >= 13, axes should be an input instead of an attribute.
|
|
gather_node.input[1] # Use 'indices' input of 'Gather'
|
|
],
|
|
outputs=if_node.output,
|
|
name=shape_node.name.replace('Shape', 'Squeeze')
|
|
)
|
|
# Replace 'Shape', 'Gather', 'Equal', 'If' with 'Squeeze'.
|
|
subgraph.node.insert(i_shape, squeeze_node)
|
|
subgraph.node.remove(shape_node)
|
|
subgraph.node.remove(gather_node)
|
|
subgraph.node.remove(equal_node)
|
|
subgraph.node.remove(if_node)
|
|
_verbose(
|
|
f'| fold nodes: [\'{shape_node.name}\', \'{gather_node.name}\', '
|
|
f'\'{equal_node.name}\', \'{if_node.name}\'] -> \'{squeeze_node.name}\'')
|
|
break
|
|
i_if += 1
|
|
else:
|
|
break
|
|
i_equal += 1
|
|
else:
|
|
break
|
|
i_gather += 1
|
|
else:
|
|
break
|
|
i_shape += 1
|
|
|
|
_graph_fold_back_to_squeeze_recursive(graph)
|
|
|
|
|
|
def graph_extract_conditioner_projections(
|
|
graph: GraphProto,
|
|
op_type: str,
|
|
weight_pattern: str,
|
|
alias_prefix: str
|
|
):
|
|
"""
|
|
Extract conditioner projection nodes out of the backbone wrapped by diffusion.
|
|
These nodes only need to be calculated once before entering the main denoising loop,
|
|
and can be reused inside the loop. This optimizes the performance of ONNX inference.
|
|
|
|
:param graph: graph to perform the operation on
|
|
:param op_type: the ONNX operator type of the conditioner projections (usually 'Conv' or 'Gemm')
|
|
:param weight_pattern: a regular expression as pattern of the conditioner projection weight keys
|
|
:param alias_prefix: add prefixes to the outputs of extracted projection nodes
|
|
"""
|
|
node_dict: Dict[str, Tuple[str, NodeProto]] = {} # key: pattern match, value: (alias, node)
|
|
|
|
def _extract_conv_nodes_recursive(subgraph: GraphProto):
|
|
to_be_removed = []
|
|
for sub_node in subgraph.node:
|
|
if sub_node.op_type == 'If':
|
|
for attr in sub_node.attribute:
|
|
branch = onnx.helper.get_attribute_value(attr)
|
|
_extract_conv_nodes_recursive(branch)
|
|
elif sub_node.op_type == 'Loop':
|
|
for attr in sub_node.attribute:
|
|
if attr.name == 'body':
|
|
body = onnx.helper.get_attribute_value(attr)
|
|
_extract_conv_nodes_recursive(body)
|
|
elif sub_node.op_type == op_type and re.match(weight_pattern, sub_node.input[1]):
|
|
# Found node to extract
|
|
cached = node_dict.get(sub_node.input[1])
|
|
if cached is None:
|
|
out_alias = f'{alias_prefix}.{len(node_dict)}'
|
|
node_dict[sub_node.input[1]] = (out_alias, sub_node)
|
|
else:
|
|
out_alias = cached[0]
|
|
out = sub_node.output[0]
|
|
# Search for nodes downstream the extracted node and match them to the renamed output.
|
|
for dep_node in subgraph.node:
|
|
for dep_idx, dep_input in enumerate(dep_node.input):
|
|
if dep_input == out:
|
|
dep_node.input.remove(out)
|
|
dep_node.input.insert(dep_idx, out_alias)
|
|
# Add the node to the remove list.
|
|
to_be_removed.append(sub_node)
|
|
[subgraph.node.remove(_n) for _n in to_be_removed]
|
|
|
|
toplevel_entry_node_idx = toplevel_entry_node = None
|
|
# Find the **last** If node in toplevel graph
|
|
for i, n in enumerate(graph.node):
|
|
if n.op_type == 'If':
|
|
toplevel_entry_node_idx = i
|
|
toplevel_entry_node = n
|
|
# If not found, find the **last** Loop node in toplevel graph
|
|
if toplevel_entry_node is None:
|
|
for i, n in enumerate(graph.node):
|
|
if n.op_type == 'Loop':
|
|
toplevel_entry_node_idx = i
|
|
toplevel_entry_node = n
|
|
if toplevel_entry_node is not None:
|
|
for a in toplevel_entry_node.attribute:
|
|
# Apply to all sub-graphs
|
|
v = onnx.helper.get_attribute_value(a)
|
|
if isinstance(v, GraphProto):
|
|
_extract_conv_nodes_recursive(v)
|
|
|
|
# Insert the extracted nodes before the first 'If' node which carries the main denoising loop.
|
|
for key in reversed(node_dict):
|
|
alias, node = node_dict[key]
|
|
# Rename output of the node.
|
|
out_name = node.output[0]
|
|
node.output.remove(node.output[0])
|
|
node.output.insert(0, alias)
|
|
# Insert node into the main graph.
|
|
graph.node.insert(toplevel_entry_node_idx, node)
|
|
# Rename value info of the output.
|
|
for v in graph.value_info:
|
|
if v.name == out_name:
|
|
v.name = alias
|
|
break
|
|
_verbose(f'| extract conditioner projection: \'{node.name}\'')
|
|
|
|
|
|
def graph_remove_unused_values(graph: GraphProto):
|
|
used_values = set()
|
|
|
|
def _record_usage_recursive(subgraph: GraphProto):
|
|
for node in subgraph.node:
|
|
# For 'If' and 'Loop' nodes, do recording recursively
|
|
if node.op_type == 'If':
|
|
for attr in node.attribute:
|
|
branch = onnx.helper.get_attribute_value(attr)
|
|
_record_usage_recursive(branch)
|
|
elif node.op_type == 'Loop':
|
|
for attr in node.attribute:
|
|
if attr.name == 'body':
|
|
body = onnx.helper.get_attribute_value(attr)
|
|
_record_usage_recursive(body)
|
|
# For each node, record its inputs and outputs
|
|
for io_list in [node.input, node.output]:
|
|
for io_value in io_list:
|
|
used_values.add(io_value)
|
|
|
|
def _clean_unused_recursively(subgraph):
|
|
# Do cleaning in sub-graphs recursively.
|
|
for node in subgraph.node:
|
|
if node.op_type == 'If':
|
|
for attr in node.attribute:
|
|
branch = onnx.helper.get_attribute_value(attr)
|
|
_clean_unused_recursively(branch)
|
|
elif node.op_type == 'Loop':
|
|
for attr in node.attribute:
|
|
if attr.name == 'body':
|
|
body = onnx.helper.get_attribute_value(attr)
|
|
_clean_unused_recursively(body)
|
|
|
|
# Do cleaning in current graph.
|
|
i = 0
|
|
while i < len(subgraph.initializer):
|
|
name = subgraph.initializer[i].name
|
|
if name not in used_values:
|
|
subgraph.initializer.pop(i)
|
|
_verbose(f'| remove unused initializer: {name}')
|
|
else:
|
|
i += 1
|
|
i = 0
|
|
while i < len(subgraph.value_info):
|
|
name = subgraph.value_info[i].name
|
|
if name not in used_values:
|
|
subgraph.value_info.pop(i)
|
|
_verbose(f'| remove unused value info: {name}')
|
|
else:
|
|
i += 1
|
|
|
|
_record_usage_recursive(graph)
|
|
_clean_unused_recursively(graph)
|