def generate_output_mask(cls, node: NNCFNode,
graph: NNCFGraph) -> Union[tf.Tensor, None]:
"""
Generate output mask from input masks with all None replaced by identity masks.
If all input masks is None return None.
:param node: Node to determine it's sources
:param graph: NNCF graph to work with
:return: Output mask
"""
previous_nodes = graph.get_previous_nodes_sorted_by_in_port(node)
input_masks = [
input_node.data['output_mask'] for input_node in previous_nodes
]
input_edges = graph.get_input_edges(node)
if all(mask is None for mask in input_masks):
return None
device = [m for m in input_masks if m is not None][0].device
filled_input_masks = []
for i, mask in enumerate(input_masks):
if mask is None:
with tf.device(device):
mask = tf.ones(input_edges[i][
NNCFGraph.ACTIVATION_SHAPE_EDGE_ATTR][-1])
filled_input_masks.append(mask)
result_mask = tf.concat(filled_input_masks, 0)
return result_mask
def test_concat_output_tensor_device():
graph = NNCFGraph()
dummy_ops = [
graph.add_nncf_node(f'dummy_op_{i}', DummyMaskProducerMetatype.name,
DummyMaskProducerMetatype) for i in range(3)
]
concat_layer_attributes = MultipleInputLayerAttributes(2)
concat_node = graph.add_nncf_node('concat_node',
'concat',
dummy_types.DummyConcatMetatype,
layer_attributes=concat_layer_attributes)
for op in dummy_ops:
graph.add_edge_between_nncf_nodes(from_node_id=op.node_id,
to_node_id=concat_node.node_id,
tensor_shape=[10] * 4,
input_port_id=0,
output_port_id=0,
dtype=Dtype.FLOAT)
# Set mask to last dummy node
ref_device = 'some_test_device'
for op in dummy_ops[:-1]:
op = graph.get_node_by_id(op.node_id)
op.data['output_mask'] = None
last_op = graph.get_node_by_id(dummy_ops[-1].node_id)
last_op.data['output_mask'] = NPNNCFTensor(np.ones(10),
dummy_device=ref_device)
# Propagate masks
MaskPropagationAlgorithm(graph,
dummy_types.DUMMY_PRUNING_OPERATOR_METATYPES,
NPNNCFTensorProcessor).mask_propagation()
# Check concat op has appropriate device
concat_node = graph.get_node_by_id(concat_node.node_id)
assert concat_node.data['output_mask'].device == ref_device
def test_conv_pruning_ops(transpose, layer_attributes, ref_accept_pruned_input,
conv_type):
default_conv_params = {
'weight_requires_grad': True,
'kernel_size': (2, 2),
'stride': (1, 1),
'padding_values': [0, 0]
}
graph = NNCFGraph()
dummy_op_before = graph.add_nncf_node('dummy_op_before',
DummyMaskProducerMetatype.name,
DummyMaskProducerMetatype)
target_conv_attributes = ConvolutionLayerAttributes(transpose=transpose,
**layer_attributes,
**default_conv_params)
conv_op_target = graph.add_nncf_node(
'conv_op_target',
dummy_types.DummyConvMetatype.name,
dummy_types.DummyConvMetatype,
layer_attributes=target_conv_attributes)
graph.add_edge_between_nncf_nodes(
from_node_id=dummy_op_before.node_id,
to_node_id=conv_op_target.node_id,
tensor_shape=[layer_attributes['in_channels']] * 4,
input_port_id=0,
output_port_id=0,
dtype=Dtype.FLOAT)
pruning_op_class = dummy_types.DummyTransposeConvPruningOp if transpose else dummy_types.DummyConvPruningOp
assert pruning_op_class.accept_pruned_input(
conv_op_target) == ref_accept_pruned_input
ones_input_mask = NPNNCFTensor(np.ones(
(layer_attributes['in_channels'], )))
ones_output_mask = NPNNCFTensor(
np.ones((layer_attributes['out_channels'], )))
# Check all combinations of masks
for input_mask in [None, ones_input_mask]:
for output_mask in [None, ones_output_mask]:
dummy_op_before = graph.get_node_by_id(dummy_op_before.node_id)
conv_op_target = graph.get_node_by_id(conv_op_target.node_id)
dummy_op_before.data['output_mask'] = input_mask
conv_op_target.data['output_mask'] = output_mask
MaskPropagationAlgorithm(
graph, dummy_types.DUMMY_PRUNING_OPERATOR_METATYPES,
NPNNCFTensorProcessor).mask_propagation()
dummy_op_before = graph.get_node_by_id(dummy_op_before.node_id)
conv_op_target = graph.get_node_by_id(conv_op_target.node_id)
if conv_type == 'usual_conv':
assert np.all(
conv_op_target.data['output_mask'] == output_mask)
elif conv_type in [
'grouped_conv_no_depthwise', 'multiply_grouped_conv'
]:
assert conv_op_target.data['output_mask'] is None
else:
assert np.all(conv_op_target.data['output_mask'] == input_mask)
def test_reshape_accept_pruned_input(node_type, input_shape, output_shape):
node_name = 'dummy_reshape'
layer_attributes = ReshapeLayerAttributes(input_shape, output_shape)
graph = NNCFGraph()
node = graph.add_nncf_node(node_name,
node_type,
METATYPES_MAP[node_type]['metatype'],
layer_attributes=layer_attributes)
actual_accept_pruned_input = METATYPES_MAP[node_type][
'ops'].accept_pruned_input(node)
assert actual_accept_pruned_input
def test_reshape_is_last_op(node_type):
node_name = 'dummy_reshape'
layer_attributes = None
graph = NNCFGraph()
prev_node = graph.add_nncf_node('prev_node',
dummy_types.DummyConvMetatype.name,
dummy_types.DummyConvMetatype)
reshape_node = graph.add_nncf_node(node_name,
node_type,
METATYPES_MAP[node_type]['metatype'],
layer_attributes=layer_attributes)
assert not METATYPES_MAP[node_type]['ops'].accept_pruned_input(
reshape_node)
graph.add_edge_between_nncf_nodes(from_node_id=prev_node.node_id,
to_node_id=reshape_node.node_id,
tensor_shape=[1, 32],
input_port_id=0,
output_port_id=0,
dtype=Dtype.FLOAT)
for output_mask in (None, NPNNCFTensor(np.ones((10, )))):
prev_node = graph.get_node_by_id(prev_node.node_id)
reshape_node = graph.get_node_by_id(reshape_node.node_id)
prev_node.data['output_mask'] = output_mask
METATYPES_MAP[node_type]['ops'].mask_propagation(
reshape_node, graph, NPNNCFTensorProcessor)
assert reshape_node.data['output_mask'] is None
def test_identity_mask_propogation_prune_ops(dummy_op_class):
assert dummy_op_class.accept_pruned_input(None)
graph = NNCFGraph()
conv_op = graph.add_nncf_node('conv_op', 'conv',
dummy_types.DummyConvMetatype)
identity_ops = []
for alias in dummy_op_class.get_all_op_aliases():
identity_op = graph.add_nncf_node(
'identity', alias, dummy_types.DummyIdentityMaskForwardMetatype)
graph.add_edge_between_nncf_nodes(from_node_id=conv_op.node_id,
to_node_id=identity_op.node_id,
tensor_shape=[10] * 4,
input_port_id=0,
output_port_id=0,
dtype=Dtype.FLOAT)
identity_ops.append(identity_op)
# Check with and without masks
for output_mask in [None, NPNNCFTensor(np.ones((10, )))]:
conv_op = graph.get_node_by_id(conv_op.node_id)
conv_op.data['output_mask'] = output_mask
MaskPropagationAlgorithm(graph,
dummy_types.DUMMY_PRUNING_OPERATOR_METATYPES,
NPNNCFTensorProcessor).mask_propagation()
for identity_op in identity_ops:
identity_op = graph.get_node_by_id(identity_op.node_id)
assert np.all(identity_op.data['output_mask'] == output_mask)
def test_group_norm_pruning_ops(num_channels, num_groups,
accept_pruned_input_ref):
graph = NNCFGraph()
conv_op = graph.add_nncf_node('conv_op', 'conv',
dummy_types.DummyConvMetatype)
group_norm_layer_attributes = GroupNormLayerAttributes(
True, num_channels=num_channels, num_groups=num_groups)
group_norm_op = graph.add_nncf_node(
'identity',
dummy_types.DummyGroupNormMetatype.name,
dummy_types.DummyGroupNormMetatype,
layer_attributes=group_norm_layer_attributes)
assert dummy_types.DummyGroupNormPruningOp.accept_pruned_input(
group_norm_op) == accept_pruned_input_ref
graph.add_edge_between_nncf_nodes(from_node_id=conv_op.node_id,
to_node_id=group_norm_op.node_id,
tensor_shape=[10] * 4,
input_port_id=0,
output_port_id=0,
dtype=Dtype.FLOAT)
# Check with and without masks
for output_mask in [None, NPNNCFTensor(np.ones((10, )))]:
conv_op = graph.get_node_by_id(conv_op.node_id)
conv_op.data['output_mask'] = output_mask
MaskPropagationAlgorithm(graph,
dummy_types.DUMMY_PRUNING_OPERATOR_METATYPES,
NPNNCFTensorProcessor).mask_propagation()
identity_op = graph.get_node_by_id(group_norm_op.node_id)
if not accept_pruned_input_ref:
output_mask = None
assert np.all(identity_op.data['output_mask'] == output_mask)
def _get_related_batchnorms(self, layer_name: str, group: NodesCluster,
graph: NNCFGraph) -> List[NNCFNode]:
"""
Returns List of batchnorm nodes related to the layer.
Note: Single node per layer for shared bactchnorm layers
"""
layer_nodes = [
node_ for node_ in group.nodes
if get_layer_identifier(node_) == layer_name
]
bn_nodes = []
bn_layer_names = []
for layer_node in layer_nodes:
for next_node in graph.get_next_nodes(layer_node):
for bn_node in graph.traverse_graph(next_node,
self._get_bn_for_node):
bn_layer_name = get_layer_identifier(bn_node)
if bn_layer_name not in bn_layer_names:
bn_layer_names.append(bn_layer_name)
bn_nodes.append(bn_node)
return bn_nodes
def _get_nncf_graph_from_sequential(model: tf.keras.Model) -> NNCFGraph:
nncf_graph = NNCFGraph()
producer_layer = None
model_config = model.get_config()
for layer in model_config['layers']:
layer_name = layer['config']['name']
layer_type = _get_layer_type(layer)
layer_dtype = _get_layer_dtype(layer)
data_format = layer['config'].get('data_format')
attrs = dict(type=layer_type,
dtype=layer_dtype,
data_format=data_format,
in_ports=[0],
out_ports=[0],
is_shared=False)
if layer_type in GENERAL_CONV_LAYERS:
module_attributes = _get_module_attributes(
model.get_layer(layer_name), attrs)
attrs.update({NNCFGraph.MODULE_ATTRIBUTES: module_attributes})
nncf_graph.add_node(layer_name, **attrs)
if producer_layer is not None:
input_shape = _prepare_shape(
model.get_layer(layer_name).input_shape)
attr = {
NNCFGraph.ACTIVATION_SHAPE_EDGE_ATTR: input_shape[0],
NNCFGraph.IN_PORT_NAME_EDGE_ATTR: 0
}
nncf_graph.add_edge(producer_layer, layer_name, **attr)
producer_layer = layer_name
return nncf_graph
def test_elementwise_prune_ops(valid_masks):
graph = NNCFGraph()
conv_op_0 = graph.add_nncf_node('conv_op_0',
dummy_types.DummyConvMetatype.name,
dummy_types.DummyConvMetatype)
conv_op_1 = graph.add_nncf_node('conv_op_1',
dummy_types.DummyConvMetatype.name,
dummy_types.DummyConvMetatype)
elementwise_op = graph.add_nncf_node(
'elementwise', dummy_types.DummyElementwiseMetatype.name,
dummy_types.DummyElementwiseMetatype)
add_node = partial(graph.add_edge_between_nncf_nodes,
tensor_shape=[10] * 4,
input_port_id=0,
output_port_id=0,
dtype=Dtype.FLOAT)
# conv_op_0 -> elementwise
add_node(from_node_id=conv_op_0.node_id, to_node_id=elementwise_op.node_id)
# conv_op_1 -> elementwise
add_node(from_node_id=conv_op_1.node_id, to_node_id=elementwise_op.node_id)
masks = [NPNNCFTensor(np.ones(
(10, ))), NPNNCFTensor(np.ones(
(10, )))] if valid_masks is not None else [None, None]
def set_masks(masks, ops):
for conv_op, mask in zip(ops, masks):
conv_op = graph.get_node_by_id(conv_op.node_id)
conv_op.data['output_mask'] = mask
if valid_masks is None or valid_masks:
if valid_masks:
set_masks(masks, [conv_op_0, conv_op_1])
MaskPropagationAlgorithm(graph,
dummy_types.DUMMY_PRUNING_OPERATOR_METATYPES,
NPNNCFTensorProcessor).mask_propagation()
elementwise_op = graph.get_node_by_id(elementwise_op.node_id)
assert np.all(elementwise_op.data['output_mask'] == masks[0])
else:
def check_wrong_masks(masks):
with pytest.raises(AssertionError):
set_masks(masks, [conv_op_0, conv_op_1])
MaskPropagationAlgorithm(
graph, dummy_types.DUMMY_PRUNING_OPERATOR_METATYPES,
NPNNCFTensorProcessor).mask_propagation()
masks[0].tensor[0] = 0
check_wrong_masks(masks)
masks[0] = NPNNCFTensorProcessor.concatenate(
[masks[1], NPNNCFTensor(np.array([1]))], axis=0)
check_wrong_masks(masks)
def test_reshape_metatype_mask_prop(node_type, input_shape, output_shape,
output_mask, output_mask_ref):
node_name = 'dummy_reshape'
layer_attributes = ReshapeLayerAttributes(input_shape, output_shape)
graph = NNCFGraph()
prev_node = graph.add_nncf_node('prev_node',
dummy_types.DummyConvMetatype.name,
dummy_types.DummyConvMetatype)
reshape_node = graph.add_nncf_node(node_name,
node_type,
METATYPES_MAP[node_type]['metatype'],
layer_attributes=layer_attributes)
graph.add_edge_between_nncf_nodes(from_node_id=prev_node.node_id,
to_node_id=reshape_node.node_id,
tensor_shape=output_shape,
input_port_id=0,
output_port_id=0,
dtype=Dtype.FLOAT)
# Check both None mask and not None mask
for output_mask_cur, output_mask_ref_cur in ([(None, None),
(output_mask,
output_mask_ref)]):
# Get reference to graph node
prev_node = graph.get_node_by_id(prev_node.node_id)
reshape_node = graph.get_node_by_id(reshape_node.node_id)
prev_node.data['output_mask'] = NPNNCFTensor(
output_mask_cur) if output_mask_cur is not None else None
if isinstance(output_mask_ref_cur, str):
with pytest.raises(AssertionError):
METATYPES_MAP[node_type]['ops'].mask_propagation(
reshape_node, graph, NPNNCFTensorProcessor)
else:
METATYPES_MAP[node_type]['ops'].mask_propagation(
reshape_node, graph, NPNNCFTensorProcessor)
if output_mask_ref_cur is None:
assert reshape_node.data['output_mask'] is None
else:
assert np.all(reshape_node.data['output_mask'].tensor ==
output_mask_ref_cur)
def check_concat(cls, node: NNCFNode, graph: NNCFGraph) -> bool:
"""
Return whether all input sources of node is convolutions or not.
:param node: Node to determine it's sources
:param graph: NNCF graph to work with
:return: True if all input sources of node is convolutions
"""
for input_node in graph.get_previous_nodes(node):
# If input has mask -> it went from convolution (source of this node is a convolution)
if input_node.data.get('output_mask', None) is None:
continue
source_nodes = get_sources_of_node(
input_node, graph,
TFConvolution.get_all_op_aliases() +
TFStopMaskForwardOps.get_all_op_aliases() +
TFInput.get_all_op_aliases())
sources_types = [node.node_type for node in source_nodes]
if any(t in sources_types
for t in TFStopMaskForwardOps.get_all_op_aliases()):
return False
return True
def test_convs_elementwise_source_before_concat(empty_mask_right_branch,
empty_mask_left_branch,
right_branch_output_channels):
graph = NNCFGraph()
conv_op_0 = graph.add_nncf_node('conv_op_0', 'conv',
dummy_types.DummyConvMetatype)
conv_op_1 = graph.add_nncf_node('conv_op_1', 'conv',
dummy_types.DummyConvMetatype)
conv_op_2 = graph.add_nncf_node('conv_op_2', 'conv',
dummy_types.DummyConvMetatype)
elementwise_node = graph.add_nncf_node(
'elementwise_node', 'elementwise',
dummy_types.DummyElementwiseMetatype)
concat_layer_attributes = MultipleInputLayerAttributes(2)
concat_node = graph.add_nncf_node('concat_node',
'concat',
dummy_types.DummyConcatMetatype,
layer_attributes=concat_layer_attributes)
add_node = partial(graph.add_edge_between_nncf_nodes,
input_port_id=0,
output_port_id=0,
dtype=Dtype.FLOAT)
# conv_op_0 -> elementwise_node
add_node(from_node_id=conv_op_0.node_id,
to_node_id=elementwise_node.node_id,
tensor_shape=[10] * 4)
# conv_op_1 -> elementwise_node
add_node(from_node_id=conv_op_1.node_id,
to_node_id=elementwise_node.node_id,
tensor_shape=[10] * 4)
# elementwise_node -> concat_node
add_node(from_node_id=elementwise_node.node_id,
to_node_id=concat_node.node_id,
tensor_shape=[10] * 4)
# conv_op_2 -> concat_node
add_node(from_node_id=conv_op_2.node_id,
to_node_id=concat_node.node_id,
tensor_shape=[10, 10, right_branch_output_channels, 10])
# Set masks
if not empty_mask_left_branch:
for conv_op in [conv_op_0, conv_op_1]:
conv_op = graph.get_node_by_id(conv_op.node_id)
conv_op.data['output_mask'] = NPNNCFTensor(np.ones(10))
if not empty_mask_right_branch:
conv_op = graph.get_node_by_id(conv_op_2.node_id)
conv_op.data['output_mask'] = NPNNCFTensor(
np.ones(right_branch_output_channels))
# Propagate masks
MaskPropagationAlgorithm(graph,
dummy_types.DUMMY_PRUNING_OPERATOR_METATYPES,
NPNNCFTensorProcessor).mask_propagation()
# Check with masks
concat_node = graph.get_node_by_id(concat_node.node_id)
if empty_mask_left_branch and empty_mask_right_branch:
assert concat_node.data['output_mask'] is None
else:
reference_mask = np.ones((10 + right_branch_output_channels, ))
np.testing.assert_equal(concat_node.data['output_mask'].tensor,
reference_mask)
def test_stop_propagate_ops(pruning_op, metatype, accept_pruned_input):
graph = NNCFGraph()
node = graph.add_nncf_node('conv_op', metatype.name, metatype)
assert pruning_op.accept_pruned_input(node) == accept_pruned_input
pruning_op.mask_propagation(node, graph, NPNNCFTensorProcessor)
assert node.data['output_mask'] is None
def _get_nncf_graph_from_raw_nodes(model_config: dict,
raw_nodes: Dict) -> NNCFGraph:
nncf_graph = NNCFGraph()
nncf_graph = _update_graph_with_raw_nodes(nncf_graph, raw_nodes,
model_config)
return nncf_graph