def change_first_conv_from_bgr_to_rgb(m):
"""For input channel format BGR model, use this function to change the first
conv weight to adapt the input into RGB.
:param m: the model proto
"""
# Check for first node.
g = m.graph
input_name = g.input[0].name
first_nodes = helper.find_following_nodes_by_input_value_name(
g, input_name)
if len(first_nodes) > 1:
return False
first_node = first_nodes[0]
# Now we have the first node. Check this first node.
if first_node.op_type != 'Conv':
return False
weight_value = helper.find_value_by_name(g, first_node.input[1])
weight_shape = helper.get_shape_from_value_info(weight_value)
if weight_shape[1] != 3:
return False
# Do weight shuffle
weight_node = helper.find_node_by_output_name(g, weight_value.name)
weight_np = helper.constant_to_numpy(weight_node)
b_channel = np.expand_dims(weight_np[:, 0, :, :], axis=1)
g_channel = np.expand_dims(weight_np[:, 1, :, :], axis=1)
r_channel = np.expand_dims(weight_np[:, 2, :, :], axis=1)
new_np = np.concatenate((r_channel, g_channel, b_channel), axis=1)
new_node = helper.numpy_to_constant(weight_value.name, new_np)
# Replace the weight and topological sort
g.node.remove(weight_node)
g.node.extend([new_node])
other.topological_sort(g)
return True
def add_bn_after(prev_node):
# Get the channel number from value info
value_name = prev_node.output[0]
value = helper.find_value_by_name(g, value_name)
shape = helper.get_shape_from_value_info(value)
channel = shape[1]
# Construct 4 weights
node_name = value_name + "_nop_bn"
ones = [1.0] * channel
zeros = [0.0] * channel
scale_node = helper.list_to_constant(node_name + "_scale", [channel], ones)
bias_node = helper.list_to_constant(node_name + "_bias", [channel], zeros)
mean_node = helper.list_to_constant(node_name + "_mean", [channel], zeros)
var_node = helper.list_to_constant(node_name + "_var", [channel], ones)
# Construct BN node
bn_node = onnx.helper.make_node(
"BatchNormalization",
[value_name,
scale_node.output[0],
bias_node.output[0],
mean_node.output[0],
var_node.output[0]],
[node_name],
name = node_name
)
# Reconnect the graph
replace_node_input(n, value_name, node_name)
# Add node to the graph
g.node.extend([bn_node, scale_node, bias_node, mean_node, var_node])
def add_nop_bn_after(g, value_names):
"""Add do-nothing BatchNormalization nodes after the given value info. It will\\
take the given names as the inputs of the new node and replace the inputs\\
of the following nodes.
:param g: the graph\\
:param value_names: a list of string which are the names of value_info.
"""
for value_name in value_names:
# Find the value first
value = helper.find_value_by_name(g, value_name)
if value is None:
value = helper.find_input_by_name(g, value_name)
if value is None:
value = helper.find_output_by_name(g, value_name)
if value is None:
print("Cannot find an value_info named {}".format(value_name))
continue
# Get the channel number from value info
shape = helper.get_shape_from_value_info(value)
channel = shape[1]
# Construct 4 weights
node_name = value_name + "_nop_bn"
ones = [1.0] * channel
zeros = [0.0] * channel
scale_node = helper.list_to_constant(node_name + "_scale", [channel],
ones)
bias_node = helper.list_to_constant(node_name + "_bias", [channel],
zeros)
mean_node = helper.list_to_constant(node_name + "_mean", [channel],
zeros)
var_node = helper.list_to_constant(node_name + "_var", [channel], ones)
# Construct BN node
bn_node = onnx.helper.make_node("BatchNormalization", [
value_name, scale_node.output[0], bias_node.output[0],
mean_node.output[0], var_node.output[0]
], [node_name],
name=node_name)
# Reconnect the graph
following_nodes = helper.find_following_nodes_by_input_value_name(
g, value_name)
if len(following_nodes) > 0:
for following_node in following_nodes:
replace_node_input(following_node, value_name, node_name)
else:
# If the node is the output, replace the output with the previous input.
new_value = onnx.helper.make_tensor_value_info(
node_name, value.type.tensor_type.elem_type, shape)
output_values = []
while len(g.output):
output_values.append(g.output.pop())
while output_values:
output_value = output_values.pop()
if output_value.name == value_name:
g.output.extend([new_value])
else:
g.output.extend([output_value])
# Add node to the graph
g.node.extend([bn_node, scale_node, bias_node, mean_node, var_node])
topological_sort(g)
def add_nop_conv_after(g, value_names):
"""Add do-nothing depthwise Conv nodes after the given value info. It will\\
take the given names as the inputs of the new node and replace the inputs\\
of the following nodes.
:param g: the graph\\
:param value_names: a list of string which are the names of value_info.
"""
for value_name in value_names:
# Find the value first
value = helper.find_value_by_name(g, value_name)
if value is None:
value = helper.find_input_by_name(g, value_name)
if value is None:
value = helper.find_output_by_name(g, value_name)
if value is None:
print("Cannot find an value_info named {}".format(value_name))
continue
# Get the channel number from value info
shape = helper.get_shape_from_value_info(value)
channel = shape[1]
# Construct 4 weights
node_name = value_name + "_nop_conv"
ones = [1.0] * channel
weight_node = helper.list_to_constant(node_name + "_weight",
[channel, 1, 1, 1], ones)
# Construct BN node
conv_node = onnx.helper.make_node("Conv",
[value_name, weight_node.output[0]],
[node_name],
name=node_name,
dilations=[1, 1],
group=channel,
kernel_shape=[1, 1],
pads=[0, 0, 0, 0],
strides=[1, 1])
# Reconnect the graph
following_nodes = helper.find_following_nodes_by_input_value_name(
g, value_name)
if len(following_nodes) > 0:
for following_node in following_nodes:
replace_node_input(following_node, value_name, node_name)
else:
# If the node is the output, replace the output with the previous input.
new_value = onnx.helper.make_tensor_value_info(
node_name, value.type.tensor_type.elem_type, shape)
output_values = []
while len(g.output):
output_values.append(g.output.pop())
while output_values:
output_value = output_values.pop()
if output_value.name == value_name:
g.output.extend([new_value])
else:
g.output.extend([output_value])
# Add node to the graph
g.node.extend([conv_node, weight_node])
topological_sort(g)
def add_0_5_to_normalized_input(m):
"""For normalized input between -0.5 ~ 0.5, add 0.5 to the input to keep it
between 0 ~ 1.
:param m: the model proto
"""
g = m.graph
if len(g.input) > 1:
print("This model has multiple inputs. Cannot normalize input.")
return
input_shape = helper.get_shape_from_value_info(g.input[0])
if len(input_shape) != 4:
print("The input shape is not BCHW. Cannot normalize input.")
return
# Construct weight
ch = input_shape[1]
weight_np = np.zeros((ch, ch, 3, 3)).astype('float32')
for i in range(ch):
weight_np[i, i, 1, 1] = 1.0
new_weight = helper.numpy_to_constant("input_norm_weight", weight_np)
# Construct bias
bias_np = np.array([0.5] * ch).astype('float32')
new_bias = helper.numpy_to_constant("input_norm_bias", bias_np)
# Construct Conv
new_conv = onnx.helper.make_node(
'Conv',
['origin_input', "input_norm_weight", "input_norm_bias"],
[g.input[0].name],
name='input_norm',
dilations=[1, 1],
kernel_shape=[3, 3],
pads=[1, 1, 1, 1],
strides=[1, 1]
)
# Construct value_infos
old_input_value = g.input.pop()
weight_value = onnx.helper.make_tensor_value_info(
'input_norm_weight',
old_input_value.type.tensor_type.elem_type,
[3, 3, 3, 3]
)
bias_value = onnx.helper.make_tensor_value_info(
'input_norm_bias',
old_input_value.type.tensor_type.elem_type,
[3]
)
# Connect the graph
new_input_value = onnx.helper.make_tensor_value_info(
'origin_input',
old_input_value.type.tensor_type.elem_type,
input_shape
)
g.input.extend([new_input_value])
g.node.extend([new_weight, new_bias, new_conv])
g.value_info.extend([weight_value, bias_value, old_input_value])
# topological sort
other.topological_sort(g)
def inference_upsample_shape(g):
"""For onnx v1.4.1+, onnx cannot inference upsample output shape. Let's\\
do it ourselves. This function only inference the next upsample without\\
output shape each time.
:param g: the graph\\
:return: True if any Upsample shape is generated. Otherwise, False.
"""
for node in g.node:
if node.op_type != 'Upsample':
continue
output_value = helper.find_value_by_name(g, node.output[0])
if output_value is None:
output_value = helper.find_output_by_name(g, node.output[0])
if output_value and helper.get_shape_from_value_info(output_value):
continue
# Get input shape
input_value = helper.find_value_by_name(g, node.input[0])
if input_value is None:
continue
#raise RuntimeError("Shape for {} has not been generated.".format(node.input[0]))
if not helper.get_shape_from_value_info(input_value):
continue
#raise RuntimeError("Shape for {} is empty.".format(node.input[0]))
input_shape = helper.get_shape_from_value_info(input_value)
# Get upsample weight
weight_node = helper.find_node_by_output_name(g, node.input[1])
weight_shape, weight = helper.constant_to_list(weight_node)
if len(input_shape) != weight_shape[0]:
raise RuntimeError(
"Unmatch input shape and weight shape: {} vs {}".format(
input_shape, weight_shape))
# Calculate shape
output_shape = list(input_shape)
for i in range(len(output_shape)):
output_shape[i] = int(input_shape[i] * weight[i])
output_value = onnx.helper.make_tensor_value_info(
node.output[0], input_value.type.tensor_type.elem_type,
output_shape)
g.value_info.extend([output_value])
return True
return False
def add_bn_on_skip_branch(g):
for n in g.node:
# Find merge node (Add)
if n.op_type != 'Add':
continue
if len(n.input) != 2:
continue
# TODO: Still need to consider more cases
# Check if skip branch exist
input_node_a = helper.find_node_by_output_name(g, n.input[0])
output_of_input_node_a = helper.find_nodes_by_input_name(
g, input_node_a.output[0])
input_node_b = helper.find_node_by_output_name(g, n.input[1])
output_of_input_node_b = helper.find_nodes_by_input_name(
g, input_node_b.output[0])
if len(output_of_input_node_a) == 1 and len(
output_of_input_node_b) == 1:
continue
if len(output_of_input_node_a) == 2:
split_node = input_node_a
elif len(output_of_input_node_b) == 2:
split_node = input_node_b
else:
continue
# Get the channel number from value info
value_name = split_node.output[0]
value = helper.find_value_by_name(g, value_name)
shape = helper.get_shape_from_value_info(value)
channel = shape[1]
# Construct 4 weights
node_name = value_name + "_nop_bn"
ones = [1.0] * channel
zeros = [0.0] * channel
scale_node = helper.list_to_constant(node_name + "_scale", [channel],
ones)
bias_node = helper.list_to_constant(node_name + "_bias", [channel],
zeros)
mean_node = helper.list_to_constant(node_name + "_mean", [channel],
zeros)
var_node = helper.list_to_constant(node_name + "_var", [channel], ones)
# Construct BN node
bn_node = onnx.helper.make_node("BatchNormalization", [
value_name, scale_node.output[0], bias_node.output[0],
mean_node.output[0], var_node.output[0]
], [node_name],
name=node_name)
# Reconnect the graph
replace_node_input(n, value_name, node_name)
# Add node to the graph
g.node.extend([bn_node, scale_node, bias_node, mean_node, var_node])
topological_sort(g)
def inference_cov_shape(g):
processed = False
for node in g.node:
if node.op_type != 'Conv':
continue
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
kernel_value_info = helper.find_value_by_name(g, node.input[1])
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
_, kernel_shape = helper.find_size_shape_from_value(kernel_value_info)
_, input_shape = helper.find_size_shape_from_value(input_value_info)
if not input_shape or not kernel_shape:
continue
strides = helper.get_attribute_by_name(node, 'strides').ints
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
def add_rgb2yynn_node(m):
"""Add a conv layer which can convert rgb to yynn input.
"""
g = m.graph
if len(g.input) > 1:
print("This model has multiple inputs. Cannot change to rgb input.")
return
input_shape = helper.get_shape_from_value_info(g.input[0])
if len(input_shape) != 4:
print("The input shape is not BCHW. Cannot normalize input.")
return
# Construct weight
ch = input_shape[1]
weight_np = np.zeros((3, 3, 4, 4)).astype('float32')
weight_np[1, 1, :3, :2] = np.array([[[[0.299],
[0.587],
[0.114]]]])
weight_np[1, 1, 3, 2:] = 1.
weight_np = np.transpose(weight_np, (3, 2, 0, 1))
new_weight = helper.numpy_to_constant("input_rgb2yynn_weight", weight_np)
# Construct conv node
new_conv = onnx.helper.make_node(
'Conv',
['new_input', "input_rgb2yynn_weight"],
[g.input[0].name],
name='input_rgba2yynn',
dilations=[1, 1],
kernel_shape=[3, 3],
pads=[1, 1, 1, 1],
strides=[1, 1]
)
# Construct value_infos
old_input_value = g.input.pop()
weight_value = onnx.helper.make_tensor_value_info(
'input_rgb2yynn_weight',
old_input_value.type.tensor_type.elem_type,
[4, 4, 3, 3]
)
# Connect the graph
new_input_value = onnx.helper.make_tensor_value_info(
'new_input',
old_input_value.type.tensor_type.elem_type,
input_shape
)
g.input.extend([new_input_value])
g.node.extend([new_weight, new_conv])
g.value_info.extend([weight_value, old_input_value])
# topological sort
other.topological_sort(g)
def change_input_from_bgr_to_rgb(m):
"""For input channel format BGR model, use this function to modify the model
to accepct RGB image.If the first node is a non-group Conv. Modify weight to
adapt the input into RGB. Otherwise create a new node.
:param m: the model proto
"""
g = m.graph
if len(g.input) > 1:
print("This model has multiple inputs. Cannot change to RGB input.")
return
input_shape = helper.get_shape_from_value_info(g.input[0])
if len(input_shape) != 4 or input_shape[1] != 3:
print("The input shape is invalid for bgr conversion.")
return
# Try change conv weight first
if change_first_conv_from_bgr_to_rgb(m):
return
# Otherwise, create a special conv node and replace the input
# Construct weight
weight_np = np.zeros((3, 3, 3, 3)).astype('float32')
weight_np[0, 2, 1, 1] = 1.0
weight_np[1, 1, 1, 1] = 1.0
weight_np[2, 0, 1, 1] = 1.0
new_weight = helper.numpy_to_constant("bgr_shuffle_weight", weight_np)
# Construct Conv
new_conv = onnx.helper.make_node(
'Conv',
['rgb_input', "bgr_shuffle_weight"],
[g.input[0].name],
name='bgr_shuffle',
dilations=[1, 1],
kernel_shape=[3, 3],
pads=[1, 1, 1, 1],
strides=[1, 1]
)
# Connect the graph
old_input_value = g.input.pop()
new_input_value = onnx.helper.make_tensor_value_info(
'rgb_input',
old_input_value.type.tensor_type.elem_type,
input_shape
)
g.input.extend([new_input_value])
g.node.extend([new_weight, new_conv])
# topological sort
other.topological_sort(g)
def inference_resize_shape(g):
for node in g.node:
if node.op_type != 'Resize':
continue
output_value = helper.find_value_by_name(g, node.output[0])
output_value = helper.find_output_by_name(
g, node.output[0]) if output_value is None else output_value
if output_value is not None:
continue
if len(node.input) == 4: # input: X, roi, scales, sizes
shape_node = helper.find_node_by_output_name(g, node.input[3])
if shape_node.op_type != 'Constant':
continue
_, shape_value = helper.constant_to_list(shape_node)
output_value = onnx.helper.make_tensor_value_info(
node.output[0], onnx.TensorProto.FLOAT,
[int(v) for v in shape_value])
g.value_info.extend([output_value])
return True
else:
# If output shape is not given, inference from scales
# Get the input shape
input_value = helper.find_value_by_name(g, node.input[0])
if input_value is None:
continue
shape_value = helper.get_shape_from_value_info(input_value)
scales_node = helper.find_node_by_output_name(g, node.input[2])
if scales_node.op_type != 'Constant':
continue
_, scales_value = helper.constant_to_list(scales_node)
for i in range(len(shape_value)):
shape_value[i] *= scales_value[i]
output_value = onnx.helper.make_tensor_value_info(
node.output[0], onnx.TensorProto.FLOAT,
[int(v) for v in shape_value])
g.value_info.extend([output_value])
return True
return False
def split_ConvTranspose(model):
"""To feed our compiler, split ConvTranspose into Upsample and Conv.
:param model: the model
"""
node_to_delete = []
# Change model properties for upsample.
if model.ir_version < 3:
print("Warning: Current model IR version is not fully supported.")
model.ir_version = 4
model.opset_import[0].version = 9
g = model.graph
# Get a Convtranspose layer
for node in g.node:
# Find a Flatten node
if node.op_type != 'ConvTranspose':
continue
# Check auto_pad
auto_pad_proto = helper.get_attribute_by_name(node, "auto_pad")
if auto_pad_proto is not None:
print("Currently not split auto_pad ConvTranspose")
continue
# Check output_shape
output_shape_proto = helper.get_attribute_by_name(node, "output_shape")
if output_shape_proto is not None:
print("Currently not split output_shape ConvTranspose")
continue
# Get input shape
input_value = helper.find_value_by_name(g, node.input[0])
if input_value is None:
input_value = helper.find_input_by_name(g, node.input[0])
if input_value is None:
print("Cannot get value info named {}.".format(node.input[0]))
exit(1)
input_shape = helper.get_shape_from_value_info(input_value)
# Get attrbutes
attr = deconv_to_conv_info_extraction(input_shape, node)
# Generate Upsample scales
upsample_output_shape = list(input_shape)
upsample_output_shape[2] = (input_shape[2] - 1) * attr["strides"][0] + 1
upsample_output_shape[3] = (input_shape[3] - 1) * attr["strides"][1] + 1
upsample_node_name = node.name + "_inner_upsample"
upsample_scale_name = upsample_node_name + "_scales"
scales_np = np.ones([4]).astype('float32')
scales_np[2] = float(upsample_output_shape[2]) / input_shape[2]
scales_np[3] = float(upsample_output_shape[3]) / input_shape[3]
scales_node = helper.numpy_to_constant(upsample_scale_name, scales_np)
# Generate a Upsample layer and an internal value info
upsample_node = onnx.helper.make_node(
"Upsample",
[node.input[0], upsample_scale_name],
[upsample_node_name],
name=upsample_node_name,
mode="zeros"
)
upsample_value_info = onnx.helper.make_tensor_value_info(
upsample_node_name,
input_value.type.tensor_type.elem_type,
upsample_output_shape
)
# Check the weight layer, it may need a transpose
if attr["group"] != input_shape[1]:
weight_node = helper.find_node_by_output_name(g, node.input[1])
weight_np = helper.constant_to_numpy(weight_node)
new_weight_np = np.transpose(weight_np, [1, 0, 2, 3])
new_weight_node = helper.numpy_to_constant(node.input[1], new_weight_np)
node_to_delete.append(weight_node)
g.node.extend([new_weight_node])
value = helper.find_value_by_name(g, node.input[1])
g.value_info.remove(value)
# Generate a Conv layer
conv_node_name = node.name + "_inner_conv"
conv_node_input = [upsample_node_name]
conv_node_input.extend(node.input[1:])
conv_node = onnx.helper.make_node(
"Conv",
conv_node_input,
[node.output[0]],
name=conv_node_name,
pads=[int(i) for i in attr["conv_pads"]],
dilations=[int(i) for i in attr["dilations"]],
group=int(attr["group"]),
kernel_shape=[int(i) for i in attr["kernel_shape"]],
strides=[int(1), int(1)]
)
# Reconnect the graph
g.node.extend([scales_node, upsample_node, conv_node])
g.value_info.extend([upsample_value_info])
node_to_delete.append(node)
# Delete useless nodes
for node in node_to_delete:
g.node.remove(node)
topological_sort(g)
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
