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