def testGroupAlignedInputOutputSlices_NoGroups(self):
self.op_slice_dict = {
self.batch_norm_op: [self.batch_norm_op_slice],
self.conv_op: [self.conv_op_slice],
self.relu_op: [self.relu_op_slice]
}
self.op_group_dict = {
self.batch_norm_op_slice: self.batch_norm_op_group,
self.conv_op_slice: self.conv_op_group,
self.relu_op_slice: self.relu_op_group
}
input_op_slices = [[self.conv_op_slice]]
output_op_slices = [[self.relu_op_slice]]
aligned_op_slice_sizes = [5]
op_handler_util.group_aligned_input_output_slices(
self.batch_norm_op, [self.conv_op], [self.relu_op],
input_op_slices, output_op_slices, aligned_op_slice_sizes,
self.mock_op_reg_manager)
self.mock_op_reg_manager.group_op_slices.assert_not_called()
self.mock_op_reg_manager.process_ops.assert_called_once_with(
[self.relu_op, self.conv_op])
self.mock_op_reg_manager.process_ops_last.assert_called_once_with(
[self.batch_norm_op])
def testGroupAlignedInputOutputSlices_InputsOutputsGrouped(self):
self.op_slice_dict = {
self.batch_norm_op: [self.batch_norm_op_slice],
self.conv_op: [self.conv_op_slice],
self.relu_op: [self.relu_op_slice]
}
self.op_group_dict = {
self.batch_norm_op_slice: self.batch_norm_op_group,
self.conv_op_slice: self.conv_op_group,
self.relu_op_slice: self.relu_op_group
}
input_op_slices = [[self.conv_op_slice]]
output_op_slices = [[self.relu_op_slice]]
aligned_op_slice_sizes = [5]
op_handler_util.group_aligned_input_output_slices(
self.batch_norm_op, [], [], input_op_slices, output_op_slices,
aligned_op_slice_sizes, self.mock_op_reg_manager)
self.mock_op_reg_manager.group_op_slices.assert_has_calls([
mock.call([self.batch_norm_op_slice, self.relu_op_slice]),
mock.call([self.batch_norm_op_slice, self.conv_op_slice],
omit_source_op_slices=[])
])
self.mock_op_reg_manager.process_ops.assert_not_called()
def assign_grouping(self, op, op_reg_manager):
"""Assign grouping to the given op and updates the manager.
Args:
op: tf.Operation to assign grouping to.
op_reg_manager: OpRegularizerManager to keep track of the grouping.
"""
# Check if all input ops have groups, or tell the manager to process them.
input_ops = op_handler_util.get_input_ops(op, op_reg_manager)
input_ops_without_group = op_handler_util.get_ops_without_groups(
input_ops, op_reg_manager)
# Check if all output ops have groups, or tell the manager to process them.
output_ops = op_handler_util.get_output_ops(op, op_reg_manager)
output_ops_without_group = op_handler_util.get_ops_without_groups(
output_ops, op_reg_manager)
# Remove non-passthrough ops from outputs ops to group with.
output_ops = op_handler_util.remove_non_passthrough_ops(
output_ops, op_reg_manager)
# Only group with ops that have the same size. Process the ops that have
# mismatched size.
input_ops_to_group, input_ops_to_process = (
op_handler_util.separate_same_size_ops(op, input_ops))
output_ops_to_group, output_ops_to_process = (
op_handler_util.separate_same_size_ops(op, output_ops))
# Remove broadcast ops.
input_ops_to_process = [input_op for input_op in input_ops_to_process
if not self._is_broadcast(input_op, op_reg_manager)]
# Also process ungrouped ops.
for input_op_without_group in input_ops_without_group:
if input_op_without_group not in input_ops_to_process:
input_ops_to_process.append(input_op_without_group)
for output_op_without_group in output_ops_without_group:
if output_op_without_group not in output_ops_to_process:
output_ops_to_process.append(output_op_without_group)
# Align op slice sizes if needed.
op_slices = op_reg_manager.get_op_slices(op)
input_op_slices = op_handler_util.get_op_slices(
input_ops_to_group, op_reg_manager)
output_op_slices = op_handler_util.get_op_slices(
output_ops_to_group, op_reg_manager)
aligned_op_slice_sizes = op_handler_util.get_aligned_op_slice_sizes(
op_slices, input_op_slices, output_op_slices)
op_handler_util.reslice_ops(input_ops_to_group + [op] + output_ops_to_group,
aligned_op_slice_sizes, op_reg_manager)
# Repopulate OpSlice data, as ops may have been resliced.
input_op_slices, output_op_slices = self._get_input_output_op_slices(
input_ops_to_group, output_ops_to_group, op_reg_manager)
# Group with inputs and outputs.
op_handler_util.group_aligned_input_output_slices(
op, input_ops_to_process, output_ops_to_process, input_op_slices,
output_op_slices, aligned_op_slice_sizes, op_reg_manager)
def assign_grouping(self, op, op_reg_manager):
"""Assign grouping to the given op and updates the manager.
Args:
op: tf.Operation to assign grouping to.
op_reg_manager: OpRegularizerManager to keep track of the grouping.
"""
# TODO(a1): Consider refactoring this duplicated logic.
# Check if all input ops have groups, or tell the manager to process them.
input_ops = op_handler_util.get_input_ops(op, op_reg_manager)
input_ops_without_group = op_handler_util.get_ops_without_groups(
input_ops, op_reg_manager)
# Check if all output ops have groups, or tell the manager to process them.
output_ops = op_handler_util.get_output_ops(op, op_reg_manager)
output_ops_without_group = op_handler_util.get_ops_without_groups(
output_ops, op_reg_manager)
# Remove non-passthrough ops from outputs ops to group with.
output_ops = op_handler_util.remove_non_passthrough_ops(
output_ops, op_reg_manager)
# Only group with output ops that have the same size. Process the ops that
# have mismatched size.
input_ops_to_group = input_ops
input_ops_to_process = input_ops_without_group
output_ops_to_group, output_ops_to_process = (
op_handler_util.separate_same_size_ops(op, output_ops))
# Also process ungrouped ops.
output_ops_to_process.extend(output_ops_without_group)
# Populate OpSlice data for all relevant ops.
concat_op_slices = op_reg_manager.get_op_slices(op)
input_op_slices, output_op_slices = self._get_input_output_op_slices(
input_ops_to_group, output_ops_to_group, op_reg_manager)
# Align op slices sizes if needed.
aligned_op_slice_sizes = op_handler_util.get_aligned_op_slice_sizes(
concat_op_slices, input_op_slices, output_op_slices)
op_handler_util.reslice_concat_ops(input_ops_to_group,
aligned_op_slice_sizes,
op_reg_manager)
op_handler_util.reslice_ops(output_ops_to_group + [op],
aligned_op_slice_sizes, op_reg_manager)
# Repopulate OpSlice data, as ops may have been resliced.
input_op_slices, output_op_slices = self._get_input_output_op_slices(
input_ops_to_group, output_ops_to_group, op_reg_manager)
# Group aligned OpSlice.
op_handler_util.group_aligned_input_output_slices(
op, input_ops_to_process, output_ops_to_process, input_op_slices,
output_op_slices, aligned_op_slice_sizes, op_reg_manager)
def assign_grouping(self, op, op_reg_manager):
"""Assign grouping to the given op and updates the manager.
Args:
op: tf.Operation to assign grouping to.
op_reg_manager: OpRegularizerManager to keep track of the grouping.
"""
concat_axis = _get_concat_op_axis(op)
# Need to figure out the rank to know if axis is last.
rank = len(op.inputs[0].shape) # Rank of the first input.
if concat_axis != -1 and concat_axis != rank - 1:
# Concat is actually grouping inputs!
handler = grouping_op_handler.GroupingOpHandler()
handler.assign_grouping(op, op_reg_manager)
return
# If concat is of the last dimension, this is a `standard` concat.
# TODO(a1): Consider refactoring this duplicated logic.
# Check if all input ops have groups, or tell the manager to process them.
input_ops = op_handler_util.get_input_ops(op, op_reg_manager)
input_ops_without_group = op_handler_util.get_ops_without_groups(
input_ops, op_reg_manager)
# Check if all output ops have groups, or tell the manager to process them.
output_ops = op_handler_util.get_output_ops(op, op_reg_manager)
output_ops_without_group = op_handler_util.get_ops_without_groups(
output_ops, op_reg_manager)
# Remove non-passthrough ops from outputs ops to group with.
output_ops = op_handler_util.remove_non_passthrough_ops(
output_ops, op_reg_manager)
# Only group with output ops that have the same size. Process the ops that
# have mismatched size.
input_ops_to_group = input_ops
input_ops_to_process = input_ops_without_group
output_ops_to_group, output_ops_to_process = (
op_handler_util.separate_same_size_ops(op, output_ops))
# Also process ungrouped ops.
output_ops_to_process.extend(output_ops_without_group)
# Populate OpSlice data for all relevant ops.
concat_op_slices = op_reg_manager.get_op_slices(op)
input_op_slices, output_op_slices = self._get_input_output_op_slices(
input_ops_to_group, output_ops_to_group, op_reg_manager)
# Align op slices sizes if needed.
aligned_op_slice_sizes = op_handler_util.get_aligned_op_slice_sizes(
concat_op_slices, input_op_slices, output_op_slices)
op_handler_util.reslice_concat_ops(input_ops_to_group,
aligned_op_slice_sizes,
op_reg_manager)
op_handler_util.reslice_ops(output_ops_to_group + [op],
aligned_op_slice_sizes, op_reg_manager)
# Repopulate OpSlice data, as ops may have been resliced.
input_op_slices, output_op_slices = self._get_input_output_op_slices(
input_ops_to_group, output_ops_to_group, op_reg_manager)
# Group aligned OpSlice.
op_handler_util.group_aligned_input_output_slices(
op, input_ops_to_process, output_ops_to_process, input_op_slices,
output_op_slices, aligned_op_slice_sizes, op_reg_manager)
def assign_grouping(self, op, op_reg_manager):
"""Assign grouping to the given op and updates the manager.
Args:
op: tf.Operation to assign grouping to.
op_reg_manager: OpRegularizerManager to keep track of the grouping.
"""
assert op.type == 'DepthwiseConv2dNative'
# Get output size.
output_size = op_handler_util.get_op_size(op)
# Get input size.
input_size = op_handler_util.get_op_size(op.inputs[0].op)
# Take depth_multiplier from size of weight tensor.
depth_multiplier = op.inputs[1].shape.as_list()[-1]
if depth_multiplier == 1:
super(DepthwiseConvolutionOpHandler, self).assign_grouping(
op, op_reg_manager)
return
# Check if all input ops have groups, or tell the manager to process them.
input_ops = op_handler_util.get_input_ops(op, op_reg_manager)
input_ops_without_group = op_handler_util.get_ops_without_groups(
input_ops, op_reg_manager)
# Check if all output ops have groups, or tell the manager to process them.
output_ops = op_handler_util.get_output_ops(op, op_reg_manager)
output_ops_without_group = op_handler_util.get_ops_without_groups(
output_ops, op_reg_manager)
# Remove non-passthrough ops from outputs ops to group with.
output_ops = op_handler_util.remove_non_passthrough_ops(
output_ops, op_reg_manager)
# Only group with ops that have the same size. Process the ops that have
# mismatched size. For the input, we hardcode that inputs[0] is a normal
# input while inputs[1] is the depthwise filter.
input_ops_to_group = [input_ops[0]]
input_ops_to_process = input_ops_without_group
output_ops_to_group, output_ops_to_process = (
op_handler_util.separate_same_size_ops(op, output_ops))
# Also process ungrouped ops.
for output_op_without_group in output_ops_without_group:
if output_op_without_group not in output_ops_to_process:
output_ops_to_process.append(output_op_without_group)
# Slice ops into individual channels. For example, consider 3 input
# channels and depth_multiplier = 2. Let the input channels be [0, 1, 2]
# and the output channels be [3, 4, 5, 6, 7, 8]. The channels should be
# individually sliced and grouped with consecutive groups of size
# depth_multiplier. Thus, this would end up grouping [0, 0, 1, 1, 2, 2] and
# [3, 4, 5, 6, 7, 8] into groups (0, 3, 4), (1, 5, 6), and (2, 7, 8).
aligned_op_slice_sizes = [1] * output_size
op_handler_util.reslice_ops(
input_ops_to_group, [1] * input_size, op_reg_manager)
op_handler_util.reslice_ops(
[op] + output_ops_to_group, aligned_op_slice_sizes, op_reg_manager)
# Rearrange OpSlice to align input and output.
input_op_slices, output_op_slices = (
self._get_depth_multiplier_input_output_op_slices(
input_ops_to_group, input_size, output_ops_to_group,
op_reg_manager, depth_multiplier))
# Group with inputs and outputs.
op_handler_util.group_aligned_input_output_slices(
op, input_ops_to_process, output_ops_to_process, input_op_slices,
output_op_slices, aligned_op_slice_sizes, op_reg_manager)