def test_multiple_outputs(self):
# - +
# / \y0 y1/ \
# x split z
# |
# y (nodes are ops; edges are going up)
g = ops.Graph()
with g.as_default():
x = array_ops.placeholder(dtypes.float32, shape=[1], name='x')
y = array_ops.placeholder(dtypes.float32, shape=[2], name='y')
y0, y1 = array_ops.split(y, num_or_size_splits=2, axis=0)
z = array_ops.placeholder(dtypes.float32, shape=[1], name='z')
math_ops.add(x, y0)
math_ops.subtract(y1, z)
y1_pattern = graph_matcher.OpTypePattern('*')
minus_pattern = graph_matcher.OpTypePattern('Sub', inputs=[y1_pattern, '*'])
matcher = graph_matcher.GraphMatcher(minus_pattern)
match_results = list(matcher.match_graph(g))
self.assertEqual(1, len(match_results))
match_result = match_results[0]
self.assertEqual(y0.op, y1.op)
self.assertEqual(match_result.get_op(y1_pattern), y1.op)
self.assertEqual(match_result.get_tensor(y1_pattern), y1)
def test_ordered_pattern(self):
# + +
# / \ / \
# x y and y x should both match when ordered inputs is False.
# Even when x and y are different operations.
g = ops.Graph()
with g.as_default():
x = array_ops.placeholder(dtypes.float32, shape=[], name='x')
y = constant_op.constant(1.0, dtype=dtypes.float32)
plus = x + y
add_pattern_a = graph_matcher.OpTypePattern(
'Add|AddV2', inputs=['Const', 'Placeholder'], ordered_inputs=False)
add_pattern_b = graph_matcher.OpTypePattern(
'Add|AddV2', inputs=['Placeholder', 'Const'], ordered_inputs=False)
add_pattern_fail = graph_matcher.OpTypePattern(
'Add|AddV2', inputs=['Const', 'Placeholder'], ordered_inputs=True)
# Both add_pattern_a and add_pattern_b should match the graph since
# ordered_input was set False.
matcher_a = graph_matcher.GraphMatcher(add_pattern_a)
self.assertEqual([
match_result.get_op(add_pattern_a)
for match_result in matcher_a.match_graph(g)
], [plus.op])
matcher_b = graph_matcher.GraphMatcher(add_pattern_b)
self.assertEqual([
match_result.get_op(add_pattern_b)
for match_result in matcher_b.match_graph(g)
], [plus.op])
# But if ordered_inputs is True, the inputs list match should fail if not
# specified in the right order.
matcher_fail = graph_matcher.GraphMatcher(add_pattern_fail)
self.assertEqual(
len([
match_result.get_op(add_pattern_fail)
for match_result in matcher_fail.match_graph(g)
]), 0)
def test_conv_layer(self):
with compat.forward_compatibility_horizon(2019, 6, 7):
g = ops.Graph()
with g.as_default():
inputs = array_ops.placeholder(dtypes.float32, shape=[8, 5, 5, 3])
with contrib_ops.arg_scope([layers.batch_norm],
fused=True,
is_training=True,
trainable=True):
return layers.convolution(
inputs,
num_outputs=16,
kernel_size=3,
stride=1,
padding='VALID',
activation_fn=nn_ops.relu,
normalizer_fn=layers.batch_norm,
normalizer_params={},
weights_initializer=initializers.xavier_initializer(),
weights_regularizer=None,
biases_initializer=init_ops.zeros_initializer(),
biases_regularizer=None,
reuse=None,
trainable=True,
scope=None)
inputs_pattern = graph_matcher.OpTypePattern('*', name='inputs')
relu_pattern = graph_matcher.OpTypePattern(
'Relu',
name='relu',
inputs=[
graph_matcher.OpTypePattern(
'FusedBatchNormV3',
inputs=[
graph_matcher.OpTypePattern(
'Conv2D', inputs=[inputs_pattern, '*']), '*', '*',
'*', '*'
])
])
matcher = graph_matcher.GraphMatcher(relu_pattern)
match_results = list(matcher.match_graph(g))
self.assertEqual(1, len(match_results))
match_result = match_results[0]
self.assertEqual(match_result.get_tensor(inputs_pattern), inputs)
self.assertEqual(match_result.get_tensor('inputs'), inputs)
def test_oneof_type_pattern(self):
# - +
# / \ / \
# x y z
g = ops.Graph()
with g.as_default():
x = array_ops.placeholder(dtypes.float32, shape=[], name='x')
y = array_ops.placeholder(dtypes.float32, shape=[], name='y')
z = array_ops.placeholder(dtypes.float32, shape=[], name='z')
plus = x + y
minus = y - z
add_or_sub_pattern = graph_matcher.OpTypePattern(
'AddV2|Add|Sub', inputs=['*', '*'])
matcher = graph_matcher.GraphMatcher(add_or_sub_pattern)
self.assertEqual([
match_result.get_op(add_or_sub_pattern)
for match_result in matcher.match_graph(g)
], [plus.op, minus.op])
def test_oneof_pattern(self):
reshape_pattern = graph_matcher.OpTypePattern('Reshape')
transpose_pattern = graph_matcher.OneofPattern([
graph_matcher.OpTypePattern(
'Transpose',
name='transpose',
inputs=[
graph_matcher.OpTypePattern(
'Slice', name='slice', inputs=[reshape_pattern, '*', '*']),
'*'
]),
graph_matcher.OpTypePattern(
'Transpose', name='transpose', inputs=[reshape_pattern, '*'])
])
matcher = graph_matcher.GraphMatcher(transpose_pattern)
g = ops.Graph()
with g.as_default():
inputs = array_ops.placeholder(dtypes.float32, shape=[6])
reshape = array_ops.reshape(inputs, [2, 3])
transpose = array_ops.transpose(reshape)
[match_result] = list(matcher.match_graph(g))
self.assertEqual(match_result.get_tensor(reshape_pattern), reshape)
self.assertEqual(match_result.get_tensor('slice'), None)
self.assertEqual(match_result.get_op('transpose'), transpose.op)
g = ops.Graph()
with g.as_default():
inputs = array_ops.placeholder(dtypes.float32, shape=[6])
reshape = array_ops.reshape(inputs, [2, 3])
slicing = array_ops.slice(reshape, [0, 0], [-1, -1])
transpose = array_ops.transpose(slicing)
[match_result] = list(matcher.match_graph(g))
self.assertEqual(match_result.get_tensor(reshape_pattern), reshape)
self.assertEqual(match_result.get_tensor('slice'), slicing)
self.assertEqual(match_result.get_op('transpose'), transpose.op)
def _FindFusedBatchNorms(graph):
"""Finds all ops and tensors related to found FusedBatchNorms.
Args:
graph: Graph to inspect.
Returns:
_FusedBatchNormMatches.
"""
input_pattern = graph_matcher.OpTypePattern('*')
# In practice, the weight pattern can match a Variable or a SpaceToBatchND
# operation that follows a variable for atrous convolutions.
weight_pattern = graph_matcher.OpTypePattern('*')
gamma_pattern = graph_matcher.OpTypePattern('*')
beta_pattern = graph_matcher.OpTypePattern('*')
mean_pattern = graph_matcher.OpTypePattern('*')
variance_pattern = graph_matcher.OpTypePattern('*')
moving_average_pattern = graph_matcher.OpTypePattern('*')
bn_decay_pattern = graph_matcher.OpTypePattern('*')
layer_pattern = graph_matcher.OpTypePattern(
'Conv2D|DepthwiseConv2dNative|MatMul',
inputs=[input_pattern, weight_pattern])
batch_to_space_pattern = graph_matcher.OpTypePattern(
'BatchToSpaceND',
inputs=[
layer_pattern,
graph_matcher.OpTypePattern('*'),
graph_matcher.OpTypePattern('*')
])
# Identity between conv/matmul and bn
layer_pattern_with_identity = graph_matcher.OpTypePattern(
'Identity',
inputs=[
graph_matcher.OneofPattern([batch_to_space_pattern, layer_pattern])
])
layer_output_pattern = graph_matcher.OneofPattern(
[layer_pattern_with_identity, layer_pattern, batch_to_space_pattern])
# MatMul has a Reshape between it and FusedBatchNorm.
matmul_reshape_pattern = graph_matcher.OpTypePattern(
'Reshape',
inputs=[layer_output_pattern,
graph_matcher.OpTypePattern('*')])
batch_norm_pattern = graph_matcher.OpTypePattern(
'FusedBatchNorm|FusedBatchNormV3',
inputs=[
graph_matcher.OneofPattern(
[matmul_reshape_pattern, layer_output_pattern]), gamma_pattern,
beta_pattern, mean_pattern, variance_pattern
])
matmul_bn_output_reshape_pattern = graph_matcher.OpTypePattern(
'Reshape',
inputs=[batch_norm_pattern,
graph_matcher.OpTypePattern('*')])
batch_norm_identity_pattern = graph_matcher.OpTypePattern(
'Identity',
inputs=[batch_norm_pattern, matmul_bn_output_reshape_pattern])
bn_identity_matcher = graph_matcher.GraphMatcher(
batch_norm_identity_pattern)
bn_matcher = graph_matcher.GraphMatcher(
graph_matcher.OneofPattern(
[matmul_bn_output_reshape_pattern, batch_norm_pattern]))
moving_average_sub_pattern = graph_matcher.OpTypePattern(
'Sub', inputs=[moving_average_pattern, batch_norm_pattern])
moving_average_mul_pattern = graph_matcher.OpTypePattern(
'Mul', inputs=[moving_average_sub_pattern, bn_decay_pattern])
moving_avg_mul_matcher = graph_matcher.GraphMatcher(
moving_average_mul_pattern)
def _GetLayerMatch(match_result):
"""Populates a layer match object containing ops/tensors for folding BNs.
Args:
match_result: Matched result from graph matcher
Returns:
layer_op: Matching conv/fc op prior to batch norm
BatchNormMatch: _BatchNormMatch containing all required batch norm
parameters.
"""
moving_mean_tensor = None
moving_variance_tensor = None
bn_decay_mean_tensor = None
bn_decay_var_tensor = None
batch_to_space_op = None
layer_op = match_result.get_op(layer_pattern)
layer_tensor = match_result.get_tensor(layer_pattern)
bn_id_op = match_result.get_op(batch_norm_identity_pattern)
bn_op = match_result.get_op(batch_norm_pattern)
if bn_id_op is None:
bn_id_op = bn_op
batch_epsilon = bn_op.get_attr('epsilon')
# In the MatMul case, the output of batch norm is reshaped back into a
# 2D tensor, so the output_tensor is the output of the Reshape op.
output_tensor = bn_op.outputs[0]
if layer_op.type == 'MatMul':
output_reshape_op = match_result.get_op(
matmul_bn_output_reshape_pattern)
# If the matcher didn't match matmul_bn_output_reshape, there will be
# another match for this 'MatMul' later, so we can skip this one.
if output_reshape_op is None:
return None, None
output_tensor = output_reshape_op.outputs[0]
# Ensure that the output tensor has consumers, otherwise this is a dangling
# node and not a match.
if not output_tensor.consumers():
return None, None
batch_to_space_op = match_result.get_op(batch_to_space_pattern)
input_tensor = match_result.get_tensor(input_pattern)
weight_tensor = match_result.get_tensor(weight_pattern)
gamma_tensor = match_result.get_tensor(gamma_pattern)
beta_tensor = match_result.get_tensor(beta_pattern)
# FusedBatchNorm in training is different from that in inference. It takes
# empty 'mean' and empty 'variance', and produces the mean and the variance
# of the batch. Therefore, when is_training is true, mean_tensor and
# variance_tensor point to 1st and 2nd (0-based) output of bn_op,
# respectively; when is_training is false, they point to bn_op's inputs.
is_training = bn_op.get_attr('is_training')
if is_training:
# FusedBatchNormGrad doesn't compute gradients of the batch_mean and
# batch_variance outputs, so we need to substitute our own custom
# gradient.
# TODO(suharshs, raghuramank): Find a way to avoid needing this hack.
# pylint: disable=protected-access
bn_op._set_attr(
'_gradient_op_type',
attr_value_pb2.AttrValue(
s=compat.as_bytes('FoldFusedBatchNormGrad')))
# pylint: enable=protected-access
mean_tensor = bn_op.outputs[1]
# The batch variance used during forward and backward prop is biased,
# i.e it is calculated as: V=sum(x(k)-mu)^2/N. For the moving average
# calculation, the variance is corrected by the term N/N-1 (Bessel's
# correction). The variance tensor read from FuseBatchNorm has Bessel's
# correction applied, so we undo it here.
scope, sep, _ = bn_op.name.rpartition('/')
g = ops.get_default_graph()
with g.as_default(), g.name_scope(scope + sep):
n = math_ops.cast(
array_ops.size(layer_tensor) / array_ops.size(mean_tensor),
dtypes.float32)
variance_tensor = math_ops.multiply(
bn_op.outputs[2], (n - 1) / n,
name='Undo_Bessel_Correction')
# TODO(suharshs): Find a way to get rid of this inner match.
for mul_match_result in moving_avg_mul_matcher.match_graph(graph):
sub_op = mul_match_result.get_op(moving_average_sub_pattern)
if sub_op.inputs[1].name == bn_op.outputs[1].name:
# During training: Batch Mean is bn_op.outputs[1]
moving_mean_tensor = sub_op.inputs[0]
bn_decay_mean_tensor = mul_match_result.get_tensor(
bn_decay_pattern)
if sub_op.inputs[1].name == bn_op.outputs[2].name:
# During training: Batch Var is bn_op.outputs[2]
moving_variance_tensor = sub_op.inputs[0]
bn_decay_var_tensor = mul_match_result.get_tensor(
bn_decay_pattern)
else:
mean_tensor = match_result.get_tensor(mean_pattern)
variance_tensor = match_result.get_tensor(variance_pattern)
return layer_op, _BatchNormMatch(
layer_op=layer_op,
bn_op=bn_op,
output_tensor=output_tensor,
input_tensor=input_tensor,
weight_tensor=weight_tensor,
gamma_tensor=gamma_tensor,
beta_tensor=beta_tensor,
mean_tensor=mean_tensor,
variance_tensor=variance_tensor,
moving_mean_tensor=moving_mean_tensor,
moving_variance_tensor=moving_variance_tensor,
bn_decay_mean_tensor=bn_decay_mean_tensor,
bn_decay_var_tensor=bn_decay_var_tensor,
batch_epsilon=batch_epsilon,
batch_to_space_op=batch_to_space_op)
layer_matches = []
# We use matched_layer_set to ensure that layers aren't matched multiple
# times.
matched_layer_set = set()
for match_result in bn_identity_matcher.match_graph(graph):
layer_op, layer_match = _GetLayerMatch(match_result)
if layer_op is not None:
if layer_op not in matched_layer_set:
matched_layer_set.add(layer_op)
layer_matches.append(layer_match)
for match_result in bn_matcher.match_graph(graph):
layer_op, layer_match = _GetLayerMatch(match_result)
if layer_op is not None:
if layer_op not in matched_layer_set:
matched_layer_set.add(layer_op)
layer_matches.append(layer_match)
return layer_matches
def _FindLayersToQuantize(graph):
"""Matches layers in graph to quantize.
The following patterns get matched. Nodes surrounded by [] will be
optionally matched:
weight|folded_weight
/
conv|fc
|
[batch_to_space_nd]
|
[post_conv_correction]
|
[biasadd|folded_bias]
|
[bypass]
|
activation
|
[post_activation_bypass]
Match replacements:
If weight|folded_weight is found, FakeQuant is added afterwards.
If bypass is found, FakeQuant is added before and after.
If activation is found, FakeQuant is added afterwards.
If post_activation_bypass is found, FakeQuant is added afterwards.
Args:
graph: Graph to perform match on.
Returns:
list of _LayerMatches.
"""
input_pattern = graph_matcher.OpTypePattern('*')
weight_var_pattern = graph_matcher.OpTypePattern('Variable|VariableV2')
weight_partition_identity_pattern = graph_matcher.OpTypePattern(
'Identity', inputs=[weight_var_pattern])
weight_partition_concat_pattern = graph_matcher.OpTypePattern(
'ConcatV2', inputs=[weight_partition_identity_pattern, '*', '*'])
weight_identity_pattern = graph_matcher.OpTypePattern(
'Identity',
inputs=[
graph_matcher.OneofPattern([
weight_partition_identity_pattern,
weight_partition_concat_pattern,
weight_var_pattern,
])
])
weight_resource_var_pattern = graph_matcher.OpTypePattern('ReadVariableOp')
folded_weight_pattern = graph_matcher.OpTypePattern('Mul')
# The weights inputs to the layer operation can either be from the Variable or
# the folded weight (Mul).
layer_pattern = graph_matcher.OpTypePattern(
'|'.join(_QUANTIZABLE_TYPES),
inputs=[
input_pattern,
graph_matcher.OneofPattern([
weight_identity_pattern, weight_resource_var_pattern,
folded_weight_pattern
])
],
ordered_inputs=False)
# For atrous convolutions a BatchToSpaceND will occur after the depthwise
# convolution.
batch_to_space_pattern = graph_matcher.OpTypePattern(
'BatchToSpaceND',
inputs=[
layer_pattern,
graph_matcher.OpTypePattern('*'),
graph_matcher.OpTypePattern('*')
])
layer_output_pattern = graph_matcher.OneofPattern(
[batch_to_space_pattern, layer_pattern])
# For separable convolutions, we are looking for a conv, followed by a conv
# with no activations between the two.
sep_conv_pattern = graph_matcher.OpTypePattern(
'|'.join(_QUANTIZABLE_TYPES),
inputs=[
graph_matcher.OneofPattern([layer_output_pattern]),
graph_matcher.OpTypePattern('*')
],
ordered_inputs=False)
folded_bias_mul_pattern = graph_matcher.OpTypePattern(
'Mul',
inputs=[graph_matcher.OpTypePattern('*'), layer_output_pattern],
ordered_inputs=False)
post_layer_op_correction_pattern = graph_matcher.OpTypePattern(
'Add|AddV2',
inputs=[folded_bias_mul_pattern,
graph_matcher.OpTypePattern('*')],
ordered_inputs=False)
folded_bias_add_pattern = graph_matcher.OpTypePattern(
'Add|AddV2',
inputs=[
post_layer_op_correction_pattern,
graph_matcher.OpTypePattern('*')
],
ordered_inputs=False)
# batch_norms with forced updates have an Identity operation at the end.
# TODO(suharshs): Find a way to easily skip extra Identity operations. The
# current issue is that doing so can often match patterns across many layers
# incorrectly.
batch_norm_identity = graph_matcher.OpTypePattern(
'Identity', inputs=[folded_bias_add_pattern])
bias_add_pattern = graph_matcher.OpTypePattern(
'Add|AddV2|BiasAdd',
inputs=[layer_output_pattern, '*'],
ordered_inputs=False)
# The bias can come from the bias add or the folded bias add.
bypass_pattern = graph_matcher.OpTypePattern(
'Add|AddV2',
inputs=[
graph_matcher.OneofPattern([
bias_add_pattern, folded_bias_add_pattern, batch_norm_identity
]), '*'
],
ordered_inputs=False)
# The input to the activation can come from bias add, fold bias add, the
# bypasses.
# TODO(suharshs): We should ideally skip Identity operations instead of
# treating them as activations.
activation_pattern = graph_matcher.OpTypePattern(
'|'.join(_ACTIVATION_TYPES) + '|Identity',
inputs=[
graph_matcher.OneofPattern([
bias_add_pattern,
folded_bias_add_pattern,
batch_norm_identity,
bypass_pattern,
layer_pattern,
])
])
post_activation_bypass_pattern = graph_matcher.OpTypePattern(
'Add|AddV2', inputs=['*', activation_pattern], ordered_inputs=False)
# The order of the following matching blocks is very important. Since matches
# aren't guaranteed to be disjoint, we structure matches from largest to
# smallest to guarantee that the largest match always wins. Additionally, we
# ensure that we don't match layers multiple times.
layer_matches = []
# We use matched_layer_set to ensure that layers aren't matched multiple
# times.
matched_layer_set = set()
# First, we match layers that have a post activation bypass. We do this first
# to ensure we don't match only the first part of this layer, missing the
# post activation bypass node.
post_activation_bypass_layer_matcher = graph_matcher.GraphMatcher(
post_activation_bypass_pattern)
for match_result in post_activation_bypass_layer_matcher.match_graph(
graph):
layer_op = match_result.get_op(layer_pattern)
weight_tensor = match_result.get_tensor(weight_identity_pattern)
if weight_tensor is None:
weight_tensor = match_result.get_tensor(
weight_resource_var_pattern)
if weight_tensor is None:
weight_tensor = match_result.get_tensor(folded_weight_pattern)
activation_op = match_result.get_op(activation_pattern)
bias_add_op = match_result.get_op(bias_add_pattern)
if bias_add_op is None:
bias_add_op = match_result.get_op(folded_bias_add_pattern)
bypass_op = match_result.get_op(bypass_pattern)
post_activation_bypass_op = match_result.get_op(
post_activation_bypass_pattern)
if layer_op not in matched_layer_set:
matched_layer_set.add(layer_op)
layer_matches.append(
_LayerMatch(layer_op, weight_tensor, activation_op, bypass_op,
post_activation_bypass_op, bias_add_op))
# Now, we match the basic layer ending at an activation. We may get duplicate
# matches from above, but we don't add them to layer_matches.
layer_matcher = graph_matcher.GraphMatcher(activation_pattern)
for match_result in layer_matcher.match_graph(graph):
layer_op = match_result.get_op(layer_pattern)
weight_tensor = match_result.get_tensor(weight_identity_pattern)
if weight_tensor is None:
weight_tensor = match_result.get_tensor(
weight_resource_var_pattern)
if weight_tensor is None:
weight_tensor = match_result.get_tensor(folded_weight_pattern)
activation_op = match_result.get_op(activation_pattern)
bias_add_op = match_result.get_op(bias_add_pattern)
if bias_add_op is None:
bias_add_op = match_result.get_op(folded_bias_add_pattern)
bypass_op = match_result.get_op(bypass_pattern)
if layer_op not in matched_layer_set:
if not _IsSkipLayer(activation_op):
matched_layer_set.add(layer_op)
layer_matches.append(
_LayerMatch(layer_op, weight_tensor, activation_op,
bypass_op, None, bias_add_op))
# Match the final layer, where there may not be an activation and instead
# the output of the final BiasAdd must be quantized. So we treat the BiasAdd
# as the 'activation_op' in the _LayerMatch, to ensure that it's output is
# quantized.
final_layer_matcher = graph_matcher.GraphMatcher(
graph_matcher.OneofPattern([bias_add_pattern,
folded_bias_add_pattern]))
for match_result in final_layer_matcher.match_graph(graph):
layer_op = match_result.get_op(layer_pattern)
weight_tensor = match_result.get_tensor(weight_identity_pattern)
if weight_tensor is None:
weight_tensor = match_result.get_tensor(
weight_resource_var_pattern)
if weight_tensor is None:
weight_tensor = match_result.get_tensor(folded_weight_pattern)
activation_op = match_result.get_op(bias_add_pattern)
if activation_op is None:
activation_op = match_result.get_op(folded_bias_add_pattern)
if layer_op not in matched_layer_set:
matched_layer_set.add(layer_op)
layer_matches.append(
_LayerMatch(layer_op, weight_tensor, activation_op, None, None,
None))
# Look for separable convolutions here
sep_conv_matcher = graph_matcher.GraphMatcher(sep_conv_pattern)
for match_result in sep_conv_matcher.match_graph(graph):
layer_op = match_result.get_op(layer_pattern)
weight_tensor = match_result.get_tensor(weight_identity_pattern)
if weight_tensor is None:
weight_tensor = match_result.get_tensor(
weight_resource_var_pattern)
activation_op = match_result.get_op(layer_pattern)
if layer_op not in matched_layer_set:
matched_layer_set.add(layer_op)
layer_matches.append(
_LayerMatch(layer_op, weight_tensor, activation_op, None, None,
None))
return layer_matches