277 lines
11 KiB
Python
277 lines
11 KiB
Python
import onnx.helper
|
|
from collections import defaultdict
|
|
|
|
from . import helper
|
|
|
|
|
|
def add_name_to_nodes(graph):
|
|
existed_node_names = set()
|
|
for node in graph.node:
|
|
if node.name:
|
|
existed_node_names.add(node.name)
|
|
for node in graph.node:
|
|
if not node.name:
|
|
if node.output[0] not in existed_node_names:
|
|
node.name = node.output[0]
|
|
else:
|
|
node.name = helper.generate_unique_name(
|
|
node.output[0], existed_node_names
|
|
)
|
|
existed_node_names.add(node.name)
|
|
if node.op_type == "Loop":
|
|
add_name_to_nodes(node.attribute[0].g)
|
|
elif node.op_type == "Scan":
|
|
attr = helper.find_attribute_by_name(node, "body")
|
|
add_name_to_nodes(attr.g)
|
|
elif node.op_type == "If":
|
|
add_name_to_nodes(node.attribute[0].g)
|
|
add_name_to_nodes(node.attribute[1].g)
|
|
elif node.op_type == "SequenceMap":
|
|
add_name_to_nodes(node.attribute[0].g)
|
|
|
|
|
|
def replace_initializer_with_constant(graph):
|
|
"""Replace initializer with Constant node.
|
|
|
|
Args:
|
|
graph (onnx.GraphProto): onnx graph.
|
|
"""
|
|
node_names = set([node.name for node in graph.node])
|
|
input_map = {i.name: i for i in graph.input}
|
|
for initializer in graph.initializer:
|
|
if initializer.name in input_map:
|
|
value_info = input_map[initializer.name]
|
|
graph.input.remove(value_info)
|
|
new_node_name = helper.generate_unique_name(initializer.name, node_names)
|
|
node = onnx.helper.make_node(
|
|
"Constant", [], [initializer.name], name=new_node_name, value=initializer
|
|
)
|
|
graph.node.insert(0, node)
|
|
node_names.add(initializer.name)
|
|
while len(graph.initializer) > 0:
|
|
graph.initializer.remove(graph.initializer[0])
|
|
|
|
|
|
def clear_descriptions(graph):
|
|
"""Clear all descriptions in the graph.
|
|
|
|
Args:
|
|
graph (onnx.GraphProto): onnx graph.
|
|
"""
|
|
graph.doc_string = ""
|
|
for node in graph.node:
|
|
node.doc_string = ""
|
|
for value_info in graph.value_info:
|
|
value_info.doc_string = ""
|
|
for value_info in graph.input:
|
|
value_info.doc_string = ""
|
|
for value_info in graph.output:
|
|
value_info.doc_string = ""
|
|
|
|
|
|
def duplicate_not_supported_shared_weights(model):
|
|
"""Duplicate shared weights in the model.
|
|
|
|
Args:
|
|
model (onnx.ModelProto): onnx model.
|
|
|
|
Returns:
|
|
onnx.ModelProto: model with duplicated shared weights.
|
|
"""
|
|
# Check if the weight is shared.
|
|
weights = {}
|
|
weights_usage = {}
|
|
weights_usage_counter = defaultdict(int)
|
|
for initializer in model.graph.initializer:
|
|
weights[initializer.name] = ('initializer', initializer)
|
|
weights_usage[initializer.name] = defaultdict(list)
|
|
for node in model.graph.node:
|
|
for input_name in node.input:
|
|
if input_name in weights:
|
|
weights_usage[input_name][node.op_type].append(node)
|
|
weights_usage_counter[input_name] += 1
|
|
if node.op_type == "Constant":
|
|
weights[node.output[0]] = ('constant', node)
|
|
weights_usage[node.output[0]] = defaultdict(list)
|
|
elif node.op_type == "ConstantOfShape":
|
|
weights[node.output[0]] = ('constant_of_shape', node)
|
|
weights_usage[node.output[0]] = defaultdict(list)
|
|
# Duplicate shared weights for each op_type usage.
|
|
for weight_name in weights:
|
|
if len(weights_usage[weight_name]) < 2:
|
|
continue
|
|
weight_type, weight = weights[weight_name]
|
|
usage = weights_usage[weight_name]
|
|
counter = 0
|
|
for op_type in usage:
|
|
if counter == 0:
|
|
counter += 1
|
|
continue
|
|
# Duplicate the weight.
|
|
if weight_type == 'initializer':
|
|
tensor = weight
|
|
elif weight_type == 'constant':
|
|
tensor = weight.attribute[0].t
|
|
elif weight_type == 'constant_of_shape':
|
|
helper.logger.warning("ConstantOfShape is not supported yet.")
|
|
counter += 1
|
|
continue
|
|
new_weight_name = weight_name + '_dup_' + str(counter)
|
|
helper.logger.debug(f"Duplicating weight {weight_name} to {new_weight_name}")
|
|
if tensor.raw_data:
|
|
new_weight = onnx.helper.make_tensor(
|
|
new_weight_name, tensor.data_type, tensor.dims, tensor.raw_data, raw=True
|
|
)
|
|
else:
|
|
new_weight = onnx.helper.make_tensor(
|
|
new_weight_name, tensor.data_type, tensor.dims, tensor.float_data
|
|
)
|
|
new_weight_node = onnx.helper.make_node(
|
|
"Constant", [], [new_weight_name], name=new_weight_name, value=new_weight
|
|
)
|
|
model.graph.node.insert(0, new_weight_node)
|
|
# Replase the weight in the nodes.
|
|
for node in usage[op_type]:
|
|
for i, input_name in enumerate(node.input):
|
|
if input_name == weight_name:
|
|
node.input[i] = new_weight_name
|
|
weights_usage_counter[weight_name] -= 1
|
|
weights_usage_counter[new_weight_name] += 1
|
|
weights_usage[new_weight_name].append(op_type)
|
|
weights_usage[weight_name] = usage[:1]
|
|
# duplicate partial shared weights
|
|
for node in model.graph.node:
|
|
found_shared = False
|
|
found_non_shared = False
|
|
for node_input in node.input:
|
|
if node_input in weights_usage_counter:
|
|
if weights_usage_counter[node_input] > 1:
|
|
found_shared = True
|
|
else:
|
|
found_non_shared = True
|
|
# Find partial shared weights.
|
|
if not (found_shared and found_non_shared):
|
|
continue
|
|
# Duplicate the shared weights.
|
|
for i, input_name in enumerate(node.input):
|
|
if input_name in weights_usage_counter and weights_usage_counter[input_name] > 1:
|
|
new_weight_name = input_name + '_dup_for_' + node.name + '_input_' + str(i)
|
|
# Duplicate the weight.
|
|
weight_type, weight = weights[input_name]
|
|
if weight_type == 'initializer':
|
|
tensor = weight
|
|
elif weight_type == 'constant':
|
|
tensor = weight.attribute[0].t
|
|
elif weight_type == 'constant_of_shape':
|
|
helper.logger.warning("ConstantOfShape is not supported yet.")
|
|
continue
|
|
helper.logger.debug(f"Duplicating weight {input_name} to {new_weight_name}")
|
|
if tensor.raw_data:
|
|
new_weight = onnx.helper.make_tensor(
|
|
new_weight_name, tensor.data_type, tensor.dims, tensor.raw_data, raw=True
|
|
)
|
|
elif len(tensor.int64_data) > 0:
|
|
new_weight = onnx.helper.make_tensor(
|
|
new_weight_name, tensor.data_type, tensor.dims, tensor.int64_data
|
|
)
|
|
elif len(tensor.float_data) > 0:
|
|
new_weight = onnx.helper.make_tensor(
|
|
new_weight_name, tensor.data_type, tensor.dims, tensor.float_data
|
|
)
|
|
else:
|
|
helper.logger.error(
|
|
f"Weight {input_name} has no data to duplicate or not supported yet."
|
|
)
|
|
exit(1)
|
|
new_weight_node = onnx.helper.make_node(
|
|
"Constant", [], [new_weight_name], name=new_weight_name, value=new_weight
|
|
)
|
|
model.graph.node.insert(0, new_weight_node)
|
|
node.input[i] = new_weight_name
|
|
weights_usage_counter[new_weight_name] += 1
|
|
weights_usage_counter[input_name] -= 1
|
|
return model
|
|
|
|
|
|
def duplicate_all_shared_weights(model):
|
|
"""Duplicate all shared weights in the model.
|
|
|
|
Args:
|
|
model (onnx.ModelProto): onnx model.
|
|
|
|
Returns:
|
|
onnx.ModelProto: model with duplicated shared weights.
|
|
"""
|
|
# Check if the weight is shared.
|
|
weights = {}
|
|
weights_usage = {}
|
|
for initializer in model.graph.initializer:
|
|
weights[initializer.name] = ('initializer', initializer)
|
|
weights_usage[initializer.name] = []
|
|
for node in model.graph.node:
|
|
for input_name in node.input:
|
|
if input_name in weights:
|
|
weights_usage[input_name].append(node)
|
|
if node.op_type == "Constant":
|
|
weights[node.output[0]] = ('constant', node)
|
|
weights_usage[node.output[0]] = []
|
|
elif node.op_type == "ConstantOfShape":
|
|
weights[node.output[0]] = ('constant_of_shape', node)
|
|
weights_usage[node.output[0]] = []
|
|
# Duplicate shared weights for each op_type usage.
|
|
for weight_name in weights:
|
|
weight_type, weight = weights[weight_name]
|
|
usages = weights_usage[weight_name]
|
|
counter = 0
|
|
for usage in usages:
|
|
if counter == 0:
|
|
counter += 1
|
|
continue
|
|
# Duplicate the weight.
|
|
if weight_type == 'initializer':
|
|
tensor = weight
|
|
elif weight_type == 'constant':
|
|
tensor = weight.attribute[0].t
|
|
elif weight_type == 'constant_of_shape':
|
|
helper.logger.warning("ConstantOfShape is not supported yet.")
|
|
counter += 1
|
|
continue
|
|
new_weight_name = weight_name + '_dup_' + str(counter)
|
|
if tensor.raw_data:
|
|
new_weight = onnx.helper.make_tensor(
|
|
new_weight_name, tensor.data_type, tensor.dims, tensor.raw_data, raw=True
|
|
)
|
|
else:
|
|
new_weight = onnx.helper.make_tensor(
|
|
new_weight_name, tensor.data_type, tensor.dims, tensor.float_data
|
|
)
|
|
new_weight_node = onnx.helper.make_node(
|
|
"Constant", [], [new_weight_name], name=new_weight_name, value=new_weight
|
|
)
|
|
model.graph.node.insert(0, new_weight_node)
|
|
helper.logger.debug(f"Duplicate weight {weight_name} to {new_weight_name}")
|
|
# Replase the weight in the nodes.
|
|
for i, input_name in enumerate(usage.input):
|
|
if input_name == weight_name:
|
|
node.input[i] = new_weight_name
|
|
break
|
|
counter += 1
|
|
return model
|
|
|
|
|
|
def run_independent_passes(model, clear_des=False):
|
|
if clear_des:
|
|
clear_descriptions(model.graph)
|
|
add_name_to_nodes(model.graph)
|
|
replace_initializer_with_constant(model.graph)
|
|
return model
|
|
|
|
|
|
def run_indenpendent_postprocessing_passes(model, duplicate_shared_weight=1):
|
|
if duplicate_shared_weight == 1:
|
|
model = duplicate_not_supported_shared_weights(model)
|
|
elif duplicate_shared_weight == 2:
|
|
model = duplicate_all_shared_weights(model)
|
|
return model
|
|
|