def fuse_consecutive_reducemean(g):
node_to_del = []
for node in g.node:
# Find consecutive ReduceMean
if node.op_type != 'ReduceMean':
continue
pre_node = helper.find_node_by_output_name(g, node.input[0])
if pre_node is None or pre_node.op_type != 'ReduceMean':
continue
# Check attributes
pre_keepdims = helper.get_var_attribute_by_name(
pre_node, 'keepdims', 'int')
pre_axes = helper.get_list_attribute_by_name(pre_node, 'axes', 'int')
cur_keepdims = helper.get_var_attribute_by_name(
node, 'keepdims', 'int')
cur_axes = helper.get_list_attribute_by_name(node, 'axes', 'int')
if pre_keepdims != 0 or cur_keepdims != 0:
continue
axes = sorted(pre_axes + cur_axes)
if axes != [2, 3]:
continue
# Merge two ReduceMean into GlobalAveragePool.
new_gap_node = onnx.helper.make_node('GlobalAveragePool',
[pre_node.input[0]],
[node.output[0] + '_intermedia'],
name=node.name + '_gap')
new_flatten_node = onnx.helper.make_node(
'Flatten', [node.output[0] + '_intermedia'], [node.output[0]],
name=node.name + '_flatten',
axis=1)
# Clean up
g.node.extend([new_gap_node, new_flatten_node])
node_to_del.extend([pre_node, node])
mid_val_info = helper.find_value_by_name(g, node.input[0])
if mid_val_info:
g.value_info.remove(mid_val_info)
while node_to_del:
node = node_to_del.pop()
g.node.remove(node)
topological_sort(g)
def replace_ReduceMean_with_GlobalAveragePool(g):
"""
Replace ReduceMean with GlobalAveragePool node when available.
If there is preceeded Transpose, check the Transpose and the ReduceMean
together. If the keep_dims is set to 0, add a Flatten.
:param g: the input graph
"""
node_to_remove = []
for node in g.node:
# Find a ReduceMean layer
if node.op_type != 'ReduceMean':
continue
# Find if it have previous Transpose and its attribute meet the need.
prev_node = helper.find_node_by_output_name(g, node.input[0])
if prev_node is not None and prev_node.op_type != 'Transpose':
prev_node = None
if prev_node is not None:
perm = helper.get_list_attribute_by_name(prev_node, 'perm', 'int')
if perm != [0, 2, 3, 1]:
prev_node = None
# Check attributes
axes = helper.get_list_attribute_by_name(node, 'axes', 'int')
keepdims = helper.get_var_attribute_by_name(node, 'keepdims', 'int')
if axes is None:
continue
if prev_node is None and axes != [2, 3]:
continue
if prev_node is not None and axes != [1, 2]:
continue
if keepdims is None:
keepdims = 1
# Replace it with GlobalAveragePool
if prev_node:
input_list = prev_node.input
else:
input_list = node.input
if keepdims == 1:
output_list = node.output
else:
output_list = [node.output[0] + '_before_flatten']
flatten_node = onnx.helper.make_node("Flatten",
output_list,
node.output,
name=node.name + "_flatten",
axis=1)
g.node.extend([flatten_node])
new_node = onnx.helper.make_node("GlobalAveragePool",
input_list,
output_list,
name=node.name)
g.node.extend([new_node])
node_to_remove.append(node)
if prev_node:
value = helper.find_value_by_name(g, prev_node.output[0])
if value:
g.value_info.remove(value)
node_to_remove.append(prev_node)
for node in node_to_remove:
g.node.remove(node)
topological_sort(g)
def deconv_to_conv_info_extraction(input_size, node_proto):
"""Extract the information needed for deconv split.
:param input_size: input shape of the deconv node.\\
:param node_proto: the deconv node proto.\\
:return: a dictionary of extracted params.
"""
attr = dict()
# Get attributes from Deconv node
attr["auto_pad"] = helper.get_var_attribute_by_name(node_proto, "auto_pad", "string")
attr["dilations"] = helper.get_list_attribute_by_name(node_proto, "dilations", "int")
attr["group"] = helper.get_var_attribute_by_name(node_proto, "group", "int")
attr["kernel_shape"] = helper.get_list_attribute_by_name(node_proto, "kernel_shape", "int")
attr["output_padding"] = helper.get_list_attribute_by_name(node_proto, "output_padding", "int")
attr["pads"] = helper.get_list_attribute_by_name(node_proto, "pads", "int")
attr["strides"] = helper.get_list_attribute_by_name(node_proto, "strides", "int")
# Get output_padding
if attr["output_padding"] is None:
if attr["auto_pad"] == "SAME_LOWER" or attr["auto_pad"] == "SAME_UPPER":
attr["output_padding"] = [attr["strides"][0] - 1, attr["strides"][1]]
else:
attr["output_padding"] = [max(attr["strides"][0] - attr["kernel_shape"][0], 0),
max(attr["strides"][1] - attr["kernel_shape"][1], 0)]
# Calculate conv_padding
if attr["auto_pad"] == "SAME_LOWER" or attr["auto_pad"] == "SAME_UPPER":
pad1_h = attr["kernel_shape"][0] - (attr["kernel_shape"][0] - 1) // 2 - 1
pad1_w = attr["kernel_shape"][1] - (attr["kernel_shape"][1] - 1) // 2 - 1
head_h = min(attr["kernel_shape"][0] // 2, (attr["output_padding"][0] + 1) // 2)
head_w = min(attr["kernel_shape"][1] // 2, (attr["output_padding"][1] + 1) // 2)
tail_h = attr["output_padding"][0] - head_h
tail_w = attr["output_padding"][1] - head_w
attr["conv_pads"] = [pad1_h + head_h, pad1_w + head_w, pad1_h + tail_h, pad1_w + tail_w]
elif attr["pads"] is not None:
sum_of_pads = sum(attr["pads"])
if sum_of_pads == 0:
# Valid padding
pad1_h = attr["kernel_shape"][0] - 0 - 1
pad1_w = attr["kernel_shape"][1] - 0 - 1
head_h = 0
head_w = 0
tail_h = attr["output_padding"][0] - head_h
tail_w = attr["output_padding"][1] - head_w
attr["conv_pads"] = [pad1_h + head_h, pad1_w + head_w, pad1_h + tail_h, pad1_w + tail_w]
else:
# Calculate output shape
tmp_output_shape = [0, 0]
tmp_output_shape[0] = attr["strides"][0] * (input_size[2] - 1) + attr["output_padding"][0] + attr["kernel_shape"][0] - attr["pads"][0] - attr["pads"][2]
tmp_output_shape[1] = attr["strides"][1] * (input_size[3] - 1) + attr["output_padding"][1] + attr["kernel_shape"][1] - attr["pads"][1] - attr["pads"][3]
# Calculate real conv output shape
tmp_center_shape = [0, 0]
tmp_center_shape[0] = (input_size[2] - 1) * attr["strides"][0] + 1
tmp_center_shape[1] = (input_size[3] - 1) * attr["strides"][1] + 1
# Calculate padding
total_padding = [0, 0]
total_padding[0] = tmp_output_shape[0] - tmp_center_shape[0] + attr["kernel_shape"][0] - 1
total_padding[1] = tmp_output_shape[1] - tmp_center_shape[1] + attr["kernel_shape"][1] - 1
if total_padding[0] < 0 or total_padding[1] < 0:
raise RuntimeError(node_proto.name + " cannot infer conv padding.")
conv_pads_ = [0] * 4
conv_pads_[0] = total_padding[0] // 2
conv_pads_[1] = total_padding[1] // 2
conv_pads_[2] = total_padding[0] - total_padding[0] // 2
conv_pads_[3] = total_padding[1] - total_padding[1] // 2
attr["conv_pads"] = conv_pads_
else:
pad1_h = attr["kernel_shape"][0] - 0 - 1
pad1_w = attr["kernel_shape"][1] - 0 - 1
head_h = 0
head_w = 0
tail_h = attr["output_padding"][0] - head_h
tail_w = attr["output_padding"][1] - head_w
attr["conv_pads"] = [pad1_h + head_h, pad1_w + head_w, pad1_h + tail_h, pad1_w + tail_w]
return attr
def inference_cov_shape(g):
processed = False
for node in g.node:
# Check for Conv output shape need to be inferrenced.
if node.op_type != 'Conv':
continue
# Input shape is not ready yet. Skip.
input_value_info = helper.find_value_by_name(g, node.input[0])
if not input_value_info:
input_value_info = helper.find_input_by_name(g, node.input[0])
if not input_value_info:
continue
_, input_shape = helper.find_size_shape_from_value(input_value_info)
if not input_shape:
continue
# Output shape is already there. Skip.
output_value_info = helper.find_value_by_name(g, node.output[0])
if not output_value_info:
output_value_info = helper.find_output_by_name(g, node.output[0])
if output_value_info and \
helper.get_shape_from_value_info(output_value_info):
continue
# Now start the inference.
# If auto_pad is set, use the auto_pad.
auto_pad = helper.get_var_attribute_by_name(node, 'auto_pad', 'string')
pads = None
if auto_pad is not None and auto_pad != 'NOTSET':
if auto_pad == 'SAME_LOWER' or auto_pad == 'SAME_UPPER':
new_output_value_info = onnx.helper.make_tensor_value_info(
node.output[0],
input_value_info.type.tensor_type.elem_type,
input_shape
)
if output_value_info:
g.value_info.remove(output_value_info)
g.value_info.extend([new_output_value_info])
processed = True
continue
elif auto_pad == 'VALID':
pads = [0, 0, 0, 0]
else:
print("Unrecognized auto_pad value: " + str(auto_pad))
exit(1)
kernel_value_info = helper.find_value_by_name(g, node.input[1])
_, kernel_shape = helper.find_size_shape_from_value(kernel_value_info)
if not input_shape or not kernel_shape:
continue
strides = helper.get_attribute_by_name(node, 'strides').ints
if not pads:
pads = helper.get_attribute_by_name(node, 'pads').ints
dilation = helper.get_attribute_by_name(node, 'dilations').ints
# Pytorch model has the case where strides only have one number
if len(strides) == 1:
return strides.append(strides[0])
if len(dilation) == 1:
return dilation.append(dilation[0])
H = math.floor((input_shape[2]+pads[0]+pads[2]-\
dilation[0]*(kernel_shape[2]-1)-1)/strides[0]+1)
W = math.floor((input_shape[3]+pads[1]+pads[3]-\
dilation[1]*(kernel_shape[3]-1)-1)/strides[1]+1)
output_shape = [input_shape[0], kernel_shape[0], H, W]
new_output_value_info = onnx.helper.make_tensor_value_info(
node.output[0],
input_value_info.type.tensor_type.elem_type,
output_shape
)
processed = True
if output_value_info:
g.value_info.remove(output_value_info)
g.value_info.extend([new_output_value_info])
return processed