def replace_pattern(self, graph: Graph, match: dict): node = match['op'] N, H, W, C = match['in_data'].shape block_size = node['block_size'] graph.remove_edge(match['in_data'].id, node.id) graph.remove_edge(node.id, match['out_data'].id) dim_6D = int64_array([0, block_size, block_size, int(C / (block_size ** 2)), H, W]) order_6D = int64_array([0, 3, 4, 1, 5, 2]) dim_4D = int64_array([0, int(H * block_size), int(W * block_size), int(C / (block_size ** 2))]) reshape_6_op = Reshape(graph, dict(name=node.id + '/Reshape_to_6D')) reshape_6_const_data = Const(graph, dict(value=dim_6D)).create_node_with_data() reshape_6_data_node = reshape_6_op.create_node_with_data([match['in_data'], reshape_6_const_data]) mark_as_correct_data_layout(reshape_6_data_node.in_node(0)) order_const_data = Const(graph, dict(value=order_6D)).create_node_with_data() transpose_op = Transpose(graph, dict(name=node.id + '/Transpose')) transpose_data_node = transpose_op.create_node_with_data([reshape_6_data_node, order_const_data]) mark_as_correct_data_layout(transpose_data_node.in_node(0)) reshape_4_op = Reshape(graph, dict(name=node.id + '/Reshape_to_4D')) reshape_4_const_data = Const(graph, dict(value=dim_4D)).create_node_with_data() reshape_4_data_node = reshape_4_op.create_node_with_data([transpose_data_node, reshape_4_const_data], data_nodes=[match['out_data']]) mark_input_as_in_correct_layout(reshape_4_data_node.in_node(0), 0) mark_output_as_in_correct_layout(reshape_4_data_node.in_node(0), 0)
def mark_squeeze_reshape_concat_before_detection_output(start_nodes: list): """ The function looks for Reshape, Concat and Squeeze ops after the 'start_nodes' with 4D output and marks them with proper attributes to infer them in original NHWC layout. This is a case of the TensorFlow Object Detection API models for the SSD heads output which produces 4D tensor with bounding box deltas. :param start_nodes: list of nodes to start search from. :return: None """ q = collections.deque() q.extend(start_nodes) while len(q) != 0: cur_node = q.popleft() if cur_node.has_valid('type'): if cur_node.soft_get('type') == 'DetectionOutput': # do not go beyond the DetectionOutput node continue # the input to Reshape comes from Convolution so it will be converted from NCHW to NHWC layout in the # InsertLayoutPropagationTransposes transformation. But the output should be kept in the original layout if cur_node.soft_get('type') == 'Reshape' and len(cur_node.out_port(0).data.get_shape()) == 4: mark_output_as_in_correct_layout(cur_node, 0) # Concat should be inferred in the original layout so the input with concatenation axis should not be # updated from NHWC to NCHW layout if cur_node.soft_get('type') == 'Concat' and len(cur_node.out_port(0).data.get_shape()) == 4: cur_node.in_port(1).__setattr__('input_permutation', None) cur_node['nchw_layout'] = True cur_node.out_node(0)['nchw_layout'] = True # Squeeze should be inferred in the original layout so the input with squeeze axis should not be updated # from NHWC to NCHW layout. The input is marked as in correct layout to prevent from inserting Transpose # from NHWC to NCHW. if cur_node.soft_get('type') == 'Squeeze' and len(cur_node.in_port(0).data.get_shape()) == 4: cur_node.in_port(1).__setattr__('input_permutation', None) mark_input_as_in_correct_layout(cur_node, 0) [q.append(port.node) for port in cur_node.out_port(0).get_destinations()]
def mark_node_as_in_correct_layout_by_in_port(self, in_port): next_in_ports = self.get_next_in_ports(in_port) in_port.__setattr__('input_permutation', None) mark_input_as_in_correct_layout(in_port.node, in_port.idx) for port in next_in_ports: mark_output_as_in_correct_layout(port.get_source().node, port.get_source().idx)
def add_convolution_to_swap_xy_coordinates(graph: Graph, input_node: Node, coordinates_size: int): """ The function add convolution node after the node 'input_node' to swap xy coordinates of the boxes produced by the node 'input_node'. It is expected that box coordinates are located in the fastest changing dimension of the 'input_node' output, i.e. the input tensor could be reshaped to [num_boxes, 4] or [num_boxes, 5]. If the size is 5, then the 0-th element for each of num_boxes blocks is not changed and element 1 is swapped with element 2, element 3 is swapped with element 4. This is the case when boxes coordinates are produced by the layer "Proposal". The exact amount of elements in each block is equal to the 'coordinates_size' parameter. :param graph: graph to operate on. :param input_node: node producing boxes coordinates. :param coordinates_size: integer value equal to 4 or 5. :return convolution node that swaps coordinates. """ # swap of input tensor with 4 or 5 numbers describing boxes are supported assert (coordinates_size in [4, 5]) input_reshape_4d_node = create_op_node_with_second_input( graph, Reshape, int64_array([-1, 1, 1, coordinates_size]), dict(name=input_node.name + '/reshape_4d'), input_node) mark_input_as_in_correct_layout(input_reshape_4d_node, 0) # do not mark second input because the reshape works in initial model layout and needs to be transformed to NCHW mark_output_as_in_correct_layout(input_reshape_4d_node, 0) if coordinates_size == 5: # zero indexed element is not box coordinate ("batch id" in case of Proposal) conv_filter_data = np.array( np.array([[[[1, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 1, 0, 0, 0], [0, 0, 0, 0, 1], [0, 0, 0, 1, 0]]]], dtype=np.float32)) else: conv_filter_data = np.array( np.array( [[[[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]]], dtype=np.float32)) conv_filter_data = np.transpose(conv_filter_data, [2, 3, 0, 1]) conv_filter_const_op = Const(graph, dict(value=conv_filter_data)) conv_filter_const_node = conv_filter_const_op.create_node( [], dict(name=input_node.name + '/weights')) conv_op = Convolution( graph, { 'bias_addable': True, 'channel_dims': np.array([3]), 'batch_dims': np.array([0]), 'input_feature_channel': 0, 'output_feature_channel': 1, 'group': 1, 'layout': 'NHWC', }) return conv_op.create_node([input_reshape_4d_node, conv_filter_const_node], dict(name=input_node.name + "/conv"))
def find_and_replace_pattern(self, graph: Graph): shape_ops = graph.get_op_nodes(op='ShapeOf') # 1. Inserting Gather to N*C format on constant shape paths for shape in shape_ops: source_port = shape.in_port(0).get_source() if is_output_data_in_correct_layout(source_port.node, source_port.idx): continue # data is already in N*C format name = shape.soft_get('name', shape.id) rank = source_port.data.get_shape().size if rank in [4, 5]: index = int64_array([0, *list(range(2, rank)), 1]) else: continue # data is layout independent gather = create_op_with_const_inputs(graph, op=Gather, port_value_dict={1: index, 2: int64_array(0)}, op_attrs={'name': name + '/GatherNCHWtoNHWC'}) shape.out_port(0).get_connection().insert_node(gather) # 2. Inserting Gather/Transpose to NC* format shape_sub_graph_end_points = self.find_shape_subgraph_endpoints([shape.out_port(0) for shape in shape_ops]) for in_port in shape_sub_graph_end_points: name = in_port.node.soft_get('name', in_port.node.id) shape = in_port.data.get_shape() should_switch_layout = not any([is_output_data_in_correct_layout(port.node, port.idx) for port in in_port.node.out_ports().values() if not port.disconnected()]) should_insert_gather = should_switch_layout and len(shape) == 1 and shape.item(0) in [4, 5] should_insert_transpose = should_switch_layout and len(shape) in [4, 5] if should_insert_gather: # we should turn input permutation off to perform it with the following gather insertion in_port.__setattr__('input_permutation', None) index = int64_array([0, shape.item(0) - 1, *list(range(1, shape.item(0) - 1))]) gather = create_op_with_const_inputs(graph, op=Gather, port_value_dict={1: index, 2: int64_array(0)}, op_attrs={'name': name + '/GatherNHWCtoNCHW'}) in_port.get_connection().insert_node(gather) elif should_insert_transpose: # we should turn input permutation off to perform it with the following transpose insertion in_port.__setattr__('input_permutation', None) order = int64_array([0, len(shape) - 1, *list(range(1, len(shape) - 1))]) transpose = create_op_with_const_inputs(graph, op=Transpose, port_value_dict={1: order}, op_attrs={'name': name + '/TransposeNHWCtoNCHW', 'override_output_shape': True}) mark_input_as_in_correct_layout(transpose, 0) mark_output_as_in_correct_layout(transpose, 0) in_port.get_connection().insert_node(transpose) else: continue # data is layout independent
def find_shape_subgraph_endpoints(self, out_ports: List[Port], visited: set = None) -> Set[Port]: """ Searches for input ports of data dependent operations starting from output ports passed to the function. Condition for data dependent operations is absence of node output value. Side action: marking the sub-graph as it is in the correct layout :param out_ports: list of output ports to start search from :param visited: set of input ports that were visited to avoid visiting them more than once :return: set of input ports of data dependent operations """ if visited is None: visited = set() deque_of_in_ports = deque() for out_port in out_ports: deque_of_in_ports.extend(out_port.get_destinations()) end_points_in_ports = set() while len(deque_of_in_ports): in_port = deque_of_in_ports.popleft() if in_port in visited: continue next_in_ports = self.get_next_in_ports(in_port) if any([port.data.get_value() is None for port in next_in_ports]): end_points_in_ports.add(in_port) else: in_port.__setattr__('input_permutation', None) mark_input_as_in_correct_layout(in_port.node, in_port.idx) for port in next_in_ports: mark_output_as_in_correct_layout(port.get_source().node, port.get_source().idx) deque_of_in_ports.extend(next_in_ports) visited.add(in_port) return end_points_in_ports
def find_and_replace_pattern(self, graph: Graph): visited = set() marked_nodes = set() condition_forward = lambda n: not InsertLayoutPropagationTranspose.is_nhwc_to_nchw_transpose_needed( n) condition_backward = lambda n: not InsertLayoutPropagationTranspose.is_nchw_to_nhwc_transpose_needed( n) for node_condition in self.op_conditions: for node in graph.get_op_nodes(): if node_condition(node): log.debug( 'Detected node "{}" as a node which should be executed in the original layout' ''.format(node.soft_get('name', node.id))) forward_visited_nodes = self.bfs([node], visited, condition_forward, True) backward_visited_nodes = self.bfs([node], visited, condition_backward, False) # find "reinterp_shape" like ops which change rank of input to 4D or 5D from smaller dimensions for back_node in backward_visited_nodes: for input_node in self.get_input_nodes(back_node): if input_node not in backward_visited_nodes and not condition_forward( input_node): marked_nodes.add(input_node) # find "reinterp_shape" like ops which change rank of input from 4D or 5D to smaller dimensions for forward_node in forward_visited_nodes: for output_node in self.get_output_nodes(forward_node): if output_node not in forward_visited_nodes and not condition_backward( output_node): marked_nodes.add(output_node) marked_nodes.update(forward_visited_nodes + backward_visited_nodes) if len(marked_nodes): log.debug( 'The following nodes will be executed in the original layout: {}' ''.format([n.soft_get('name', n.id) for n in marked_nodes])) # mark all matched nodes as in correct layout and disable attributes permutation for them for visited_node in marked_nodes: mark_as_correct_data_layout(visited_node) visited_node['nchw_layout'] = True _, nodes_weigths, nodes_in_weights = self.get_ports_and_nodes_on_weights( graph) for node in nodes_weigths: if node in nodes_in_weights: for ind, port in node.in_ports().items(): if ind not in nodes_in_weights[node]: mark_input_as_in_correct_layout(node, ind) for ind, port in node.out_ports().items(): mark_output_as_in_correct_layout(node, ind) else: mark_as_correct_data_layout(node) node['nchw_layout'] = True if node.soft_get( 'type' ) == 'Const': # WA for Const op deletion during clean_up node.out_node()['nchw_layout'] = True for node in self.get_ports_and_nodes_on_shape_subgraphs(graph)[1]: mark_as_correct_data_layout(node) node['nchw_layout'] = True if node.soft_get( 'type' ) == 'Const': # WA for Const op deletion during clean_up node.out_node()['nchw_layout'] = True