def _split_prelu(node: util.NodeDef, input_node_map: util.NameToNode,
weight_modifiers: util.WeightModifiers) -> util.NodeList:
# Prelu activation is not supported by TF so this functions generates an
# equivalent formulation:
# f(x) = alpha*relu(x) if x < 0, relu(x) if x >= 0
# = pos(x)+neg(x), where pos(x)=relu(x), and neg(x)=-alpha*relu(-x)
#
# We return the sub-graph
# [pos=Relu(x), Neg(x), Relu(-x), neg=Mul(-alpha, -x), Add(pos, neg)]
inputs = list(node.input)
def _get_name(suffix):
return generate_name_from(node.name, input_node_map, suffix=suffix)
# here we need to manually keep node names unique in the sub-graph
# since we cannot modify input_node_map, because we don't have a node yet
pos = util.make_op_node('Relu', inputs[0], _get_name('Relu'))
neg_x = util.make_op_node('Neg', inputs[0], _get_name('Neg'))
neg_relu = util.make_op_node('Relu', neg_x, _get_name('Relu_1'))
neg_alpha = inputs[1]
neg = util.make_op_node('Mul', [neg_alpha, neg_relu], _get_name('Mul'))
add = util.make_op_node('Add', [pos, neg], _get_name('Add'))
# convert alpha to -alpha by registering a weight modifier function
weight_modifiers[neg_alpha] = lambda tensor: -tensor
return [pos, neg_x, neg_relu, neg, add]
def test_generate_name_from_given_duplicate_name(self):
"""generate_name_from should return new name if base name is taken"""
node_map = {}
# case #1: name without slashes -> return name + count
base_name = 'name_without_slashes'
node_map[base_name] = rewrite.make_op_node('Neg', 'x', name='_')
name = rewrite.generate_name_from(base_name, node_map)
self.assertEqual(name, base_name + '_1')
# case #2: count needs to increment
node_map[name] = rewrite.make_op_node('Neg', 'y', name='_')
name = rewrite.generate_name_from(base_name, node_map)
self.assertEqual(name, base_name + '_2')
# case #3: name + suffix -> return name + suffix + count
suffix = 'suffix'
node_map[base_name + '/' + suffix] = rewrite.make_op_node('Inv',
'x',
name='_')
name = rewrite.generate_name_from(base_name, node_map, suffix=suffix)
self.assertEqual(name, f'{base_name}/{suffix}_1')
# case #4: count needs to increment
node_map[name] = rewrite.make_op_node('Inv', 'y', name='_')
name = rewrite.generate_name_from(base_name, node_map, suffix=suffix)
self.assertEqual(name, f'{base_name}/{suffix}_2')
# case #5: slashes in name, name in map -> return name + count
base_name = 'name_without_slashes/suffix'
name = rewrite.generate_name_from(base_name, node_map)
self.assertEqual(name, 'name_without_slashes_2')
# case #6: slashes in name, name + suffix in map -> name+suffix+count
base_name = 'name_without_slashes/suffix'
name = rewrite.generate_name_from(base_name, node_map, suffix=suffix)
self.assertEqual(name, f'{base_name}_2')
def test_make_op_node_given_mixed_inputs(self):
"""make_op_node should support list of mixed types for inputs"""
node_a = rewrite.make_op_node('Neg', ['x'], name='node_a')
node = rewrite.make_op_node('Add', [node_a, 'node_b'])
self.assertEqual(node.name, 'Add')
self.assertEqual(node.op, 'Add')
self.assertEqual(node.input, ['node_a', 'node_b'])
def test_get_input_node_map_given_duplicates(self):
"""get_input_node_map should raise ValueError given duplicate names"""
graph_def = testutils.get_sample_graph_def()
relu = _get_node_by_name(graph_def, 'model/conv3/Relu')
neg = rewrite.make_op_node('Neg', list(relu.input), name='kate')
dup = rewrite.make_op_node('Exp', neg, name='model/conv3/BiasAdd')
replace_nodes = {
'model/conv3/Relu': [neg, dup],
}
updated_graph = rewrite.update_graph_def(graph_def, replace_nodes, {})
self.assertRaises(ValueError,
lambda: rewrite.get_input_node_map(updated_graph))
def _split_fused_op(node: util.NodeDef, input_node_map: util.NameToNode,
_) -> util.NodeList:
# Possible fused Conv2D patterns are:
# • Conv2D + BiasAdd
# • Conv2D + BiasAdd + <Activation>
# • Conv2D + FusedBatchNorm + <Activation>
# • Conv2D + Squeeze + BiasAdd
#
# Fused MatMul only has one pattern:
# • MatMul + BiasAdd + <Activation>
#
# FusedBatchNorm and Squeeze are not relevant for inference and thus never
# present in (optimised) frozen graphs generated by tfjs converter.
# This leaves us with Conv2D|MatMul + BiasAdd + <Activation> as the only
# remaining possible variants.
#
# For compatibility reasons with quantised TFLite models, we optionally
# split Conv2D + BiasAdd as well.
#
# We return [Conv2D|MatMul, BiasAdd|BiasAddV1, <Activation>].
# Unsupported <Activation>-nodes will be dealt with in a separate step
fused_op_name = node.op[6:] # remove the '_Fused'-prefix
fused_ops = list(s.decode('utf-8') for s in node.attr['fused_ops'].list.s)
inputs = list(node.input)
names_used = set()
def node_name(node_index):
name = generate_name_from(inputs[node_index], input_node_map)
if name in names_used:
name = generate_name_from(name,
input_node_map,
suffix=fused_ops[node_index - 2])
names_used.add(name)
return name
fused_op = util.make_op_node(fused_op_name, inputs[0:2], node_name(1))
fused_op = util.copy_op_attrs(source=node, target=fused_op)
bias_add = util.make_op_node(fused_ops[0], [fused_op, inputs[2]],
node_name(2))
bias_add = util.copy_op_attrs(source=node, target=bias_add)
have_activation = len(fused_ops) > 1
if have_activation:
# Prelu activation has additional input; reuse original name in all
# cases
name = node_name(3) if len(inputs) > 3 else node.name
activation = util.make_op_node(fused_ops[1], [bias_add] + inputs[3:],
name)
return [fused_op, bias_add, activation]
else:
return [fused_op, bias_add]
def test_update_graph_def_given_replaced_nodes(self):
"""update_graph_def should replace nodes mapped to new sub-graph"""
graph_def = testutils.get_sample_graph_def()
# let's replace the conv1 activation with log-sigmoid
relu = _get_node_by_name(graph_def, 'model/conv1/Relu')
neg = rewrite.make_op_node('Neg', list(relu.input), 'model/conv1/Neg')
exp = rewrite.make_op_node('Exp', neg, 'model/conv1/Exp')
add = rewrite.make_op_node('Add', [exp, 'one'], 'model/conv1/Add')
inv = rewrite.make_op_node('Inv', add, 'model/conv1/Inv')
replace_nodes = {'model/conv1/Relu': [neg, exp, add, inv]}
updated_graph = rewrite.update_graph_def(graph_def, replace_nodes, {})
for node_name in replace_nodes.keys():
self.assertIsNone(_get_node_by_name(updated_graph, node_name))
for node in list(replace_nodes.values())[0]:
self.assertIsNotNone(_get_node_by_name(updated_graph, node.name))
def test_copy_op_attrs(self):
"""copy_op_attrs should only copy attrs supported by the target node"""
# copy_op_attrs is used to transfer attrs from a fused op node
# (e.g. _FusedConv2D) to a standalone op (e.g. Conv2D)
# any additional attrs of the fused op need to be ignored
fused_op_str = '{"name":"model/conv2d/BiasAdd",'\
+ '"op":"_FusedConv2D","input":["input",'\
+ '"model/conv2d/Conv2D/ReadVariableOp",'\
+ '"model/conv2d/BiasAdd/ReadVariableOp",'\
+ '"model/p_re_lu/Neg"],"device":"/device:CPU:0",' \
+ '"attr":{"dilations":{"list":{"i":["1","1","1","1"]}},'\
+ '"T":{"type":"DT_FLOAT"},"data_format":{"s":"TkhXQw=="},'\
+ '"strides":{"list":{"i":["1","1","1","1"]}},'\
+ '"use_cudnn_on_gpu":{"b":true},'\
+ '"explicit_paddings":{"list":{}},'\
+ '"num_args":{"i":"2"},"epsilon":{"f":0},'\
+ '"padding":{"s":"VkFMSUQ="},'\
+ '"fused_ops":{"list":{"s":["Qmlhc0FkZA==","UHJlbHU="]}}}}'
fused_op = testutils.node_proto_from_json(fused_op_str)
node = rewrite.make_op_node('Conv2D', fused_op.input[0:2])
rewrite.copy_op_attrs(source=fused_op, target=node)
op_def = rewrite.get_op_def(node.op)
allowed = set(attr.name for attr in op_def.attr)
forbidden = any(attr for attr in node.attr if attr not in allowed)
self.assertFalse(forbidden)
# randomply check for some of the expected attributes
self.assertTrue('padding' in node.attr)
self.assertTrue('strides' in node.attr)
def rename_output_nodes(graph_def: GraphDef, name_mapping: dict) -> GraphDef:
"""Rename output nodes - the update is performed in-place.
Args:
graph_def: GraphDef proto containing the model
name_mapping: dict that maps output names to new names
Returns:
Updated GraphDef proto - same as input since the update is
performed in-place.
"""
output_nodes = get_output_nodes(graph_def)
_validate_mapping(graph_def, output_nodes, name_mapping)
nodes_to_rename = name_mapping.keys()
target_nodes = filter(lambda n: n.name in nodes_to_rename, graph_def.node)
for node in target_nodes:
new_name = name_mapping[node.name]
if node.op == 'Identity':
# dummy nodes can just be renamed
node.name = new_name
else:
# append an identity node that takes the original output as input
identity = rewrite.make_op_node('Identity',
node,
name=new_name,
dtype=_dtype(node))
graph_def.node.append(identity)
return graph_def
def test_make_op_node_accepts_dtype(self):
"""make_op_node should accept dtype parameter and default to float32"""
def dtype(node):
return rewrite.dtypes.as_dtype(node.attr['T'].type)
# from default
node = rewrite.make_op_node('Neg', 'x')
self.assertEqual(dtype(node), rewrite.dtypes.float32)
# from dtype
node = rewrite.make_op_node('Neg', 'x', dtype=rewrite.dtypes.float16)
self.assertEqual(dtype(node), rewrite.dtypes.float16)
# from enum
node = rewrite.make_op_node('Neg', 'x', dtype=6)
self.assertEqual(dtype(node), rewrite.dtypes.int8)
# from string
node = rewrite.make_op_node('Neg', 'x', dtype='half')
self.assertEqual(dtype(node), rewrite.dtypes.float16)
def _split_fused_op(node: util.NodeDef, input_node_map: util.NameToNode,
_) -> util.NodeList:
# Possible fused Conv2D patterns are:
# • Conv2D + BiasAdd
# • Conv2D + BiasAdd + <Activation>
# • Conv2D + FusedBatchNorm + <Activation>
# • Conv2D + Squeeze + BiasAdd
#
# Fused MatMul only has one pattern:
# • MatMul + BiasAdd + <Activation>
#
# FusedBatchNorm and Squeeze are not relevant for inference and thus never
# present in (optimised) frozen graphs generated by tfjs converter.
# This leaves us with Conv2D|MatMul + BiasAdd + <Activation> as the only
# remaining possible variants, since Conv2D + BiasAdd doesn't need to be
# split.
#
# We return [Conv2D|MatMul, BiasAdd|BiasAddV1, <Activation>].
# Unsupported <Activation>-nodes will be dealt with in a separate step
fused_op_name = node.op[6:] # remove the '_Fused'-prefix
fused_ops = list(s.decode('utf-8') for s in node.attr['fused_ops'].list.s)
inputs = list(node.input)
names_used = set()
def node_name(i):
name = util.generate_name_from(inputs[i], input_node_map)
if name in names_used:
# avoid name collisions by adding the name of the fused operation
# (the first name is always unique)
name = util.generate_name_from(name,
input_node_map,
suffix=fused_ops[i - 2])
names_used.add(name)
return name
fused_op = util.make_op_node(fused_op_name, inputs[0:2], node_name(1))
fused_op = util.copy_op_attrs(source=node, target=fused_op)
bias_add = util.make_op_node(fused_ops[0], [fused_op, inputs[2]],
node_name(2))
activation = util.make_op_node(fused_ops[1], [bias_add] + inputs[3:],
node_name(3))
return [fused_op, bias_add, activation]
def test_generate_name_from_given_unique_name(self):
"""generate_name_from should return base name if name is unique"""
# case #1: no slashes in name, name is unique -> return as-is
node_map = {}
base_name = 'name_without_slahes'
name = rewrite.generate_name_from(base_name, node_map)
self.assertEqual(name, base_name)
# case #2: no slashes in name, name in map, name + suffix unique
node_map[base_name] = rewrite.make_op_node('Neg', 'x', name=base_name)
suffix = 'uniq'
name = rewrite.generate_name_from(base_name, node_map, suffix=suffix)
self.assertEqual(name, f'{base_name}/{suffix}')
# case #3: slashes in name, name with slashes in map -> return 1st part
base_name = 'name/suffix'
node_map[base_name] = rewrite.make_op_node('Neg', 'y', name=base_name)
name = rewrite.generate_name_from(base_name, node_map)
self.assertEqual(name, 'name')
# case #4: slashes in name, name in map, name + suffix unique
base_name = 'name/suffix/op'
name = rewrite.generate_name_from(base_name, node_map, suffix='uniq')
self.assertEqual(name, 'name/suffix/uniq')
def _split_fused_depthwise(node: util.NodeDef, input_node_map: util.NameToNode,
weight_mods: util.WeightModifiers) -> util.NodeList:
"""Decompose fused op into DepthwiseConv2dNative + BiasAdd [+ Activation]
"""
fused_ops = list(s.decode('utf-8') for s in node.attr['fused_ops'].list.s)
inputs = node.input
names_used = set()
def node_name(node_index):
"""Return unique node names for sub-operations by appending fused-op"""
i = min(node_index, len(inputs)-1) # PReLU has 4 inputs, others only 3
name = generate_name_from(inputs[i], input_node_map)
if name in names_used:
name = generate_name_from(name, input_node_map,
suffix=fused_ops[node_index-2])
names_used.add(name)
return name
op = 'DepthwiseConv2dNative'
depthwise = util.make_op_node(op, inputs[0:2], node_name(1))
depthwise = util.copy_op_attrs(source=node, target=depthwise)
op = fused_ops[0]
bias_add = util.make_op_node(op, [depthwise, inputs[2]], node_name(2))
bias_add = util.copy_op_attrs(source=node, target=bias_add)
node_list = [depthwise, bias_add]
if len(fused_ops) > 1:
# we have an activation function
op = fused_ops[1]
input_nodes = [bias_add] + inputs[3:]
if util.get_op_def(op) is None:
# unsupported activation function - just copy type attribute
dtype = depthwise.attr['T'].type
activation = util.make_op_node(op, input_nodes, node_name(3),
dtype)
else:
# supported activation function - copy applicable attributes
activation = util.make_op_node(op, input_nodes, node_name(3))
activation = util.copy_op_attrs(source=node, target=activation)
node_list.append(activation)
return node_list
def _convert_to_log_sigmoid(node, input_map, modifiers):
"""replace Relu with logarithmic sigmoid 1/(1+exp(-x))"""
def _get_name(suffix):
return rewrite.generate_name_from(node.name, input_map,
f'logSigmoid/{suffix}')
nonlocal new_name_of_replaced_node
# -x
neg = rewrite.make_op_node('Neg',
list(node.input),
name=_get_name('Neg'))
# exp(-x)
exp = rewrite.make_op_node('Exp', neg, name=_get_name('Exp'))
# constant tensor holding "1"
res = rewrite.make_const_node(np.array([1], dtype=np.float32),
name=_get_name('Var/resource'))
# variable holding "1"
one = rewrite.make_op_node('Identity', res, _get_name('Var'))
# 1+exp(-x)
add = rewrite.make_op_node('Add', [one, exp], _get_name('Add'))
# 1/(1+exp-x)
inv = rewrite.make_op_node('Inv', add, _get_name('Inv'))
new_name_of_replaced_node = inv.name # remember the output name
return [neg, exp, res, one, add, inv]
def test_make_op_node_given_input_as_scalar(self):
"""make_op_node should accept scalar input of type string and node"""
node_from_str = rewrite.make_op_node('Neg', 'x', name='node_a')
self.assertEqual(node_from_str.input, ['x'])
node_from_node = rewrite.make_op_node('Neg', node_from_str)
self.assertEqual(node_from_node.input, ['node_a'])
def test_make_op_node_given_str_inputs(self):
"""make_op_node should support list of strings for inputs"""
node = rewrite.make_op_node('Add', ['x', 'y'])
self.assertEqual(node.name, 'Add')
self.assertEqual(node.op, 'Add')
self.assertEqual(node.input, ['x', 'y'])
