def replace_pattern(self, graph: Graph, match: dict):
bias_add = match['BiasAdd']
# Replace BiasAdd by Add operation
new_add = Add(graph, {'name': bias_add.id + '/Add'}).create_node()
bias_add.in_port(0).get_connection().set_destination(
new_add.in_port(0))
bias_add.in_port(1).get_connection().set_destination(
new_add.in_port(1))
bias_add.out_port(0).get_connection().set_source(new_add.out_port(0))
if bias_add.data_format != 'NCHW':
return
input_shape = new_add.in_port(0).data.get_shape()
bias_shape = new_add.in_port(1).data.get_shape()
assert len(bias_shape) == 1
unsqueeze_dims = np.arange(len(input_shape))
channel_dim = get_features_dim('NCHW', len(input_shape))
unsqueeze_dims = np.delete(unsqueeze_dims, channel_dim, 0)
unsqueeze_node = Unsqueeze(graph, {
'name': new_add.id + '/BiasUnsqueeze'
}).create_node()
unsqueeze_dims_node = Const(graph, {
'name': new_add.id + '/Dims',
'value': unsqueeze_dims
}).create_node()
# Reconnecting nodes
unsqueeze_node.in_port(1).connect(unsqueeze_dims_node.out_port(0))
unsqueeze_node['override_output_shape'] = True
new_add.in_port(1).get_connection().insert_node(unsqueeze_node)
def unsqueeze_num_directions(graph: Graph, match: dict):
""" Assuming considered LSTM/GRU/RNN node should has num_directions in output shape and add Unsqueeze
to match it.
"""
rnn_layer = match['rnn_layer']
rnn_layer_name = rnn_layer.soft_get('name', rnn_layer.id)
# num_directions is at 1st position in output shape, and in 0st position in hidden and cell states
# please refer to docs in this transform
direction_dim = [1, 0, 0] # index of dimension with direction index
for i in rnn_layer.out_nodes():
old_data_node = rnn_layer.out_node(i)
old_shape = old_data_node.shape.copy()
new_shape = np.delete(old_shape, direction_dim[i])
data = Op._create_data_node(graph, name=rnn_layer.name + '/Out/{}/'.format(i), attrs={'shape': new_shape})
graph.remove_edge(rnn_layer.id, old_data_node.id)
graph.add_edge(rnn_layer.id, data.id, key=0, out=i)
unsqueeze = Unsqueeze(graph, dict())
unsqueeze_dim_data = Const(graph, {'name': rnn_layer.name + '/UnsqueezeNumDirections/{}/Dim'.format(i),
'value': int64_array([direction_dim[i]])}).create_node_with_data()
unsqueeze.create_node_with_data([data, unsqueeze_dim_data],
dict(name=rnn_layer_name + '/UnsqueezeNumDirections/{}'.format(i)),
data_nodes=[old_data_node])
def test_unsqueeze_infer_negative_indices(self):
unsq_dims = np.array([-1])
graph = build_graph(self.nodes_attributes,
[('data_1', 'unsq'),
('unsq_dims_const', 'unsq_dims'),
('unsq_dims', 'unsq'),
('unsq', 'data_2')],
{'data_1': {'shape': np.array([2, 3, 64, 64])},
'unsq_dims': {'value': unsq_dims, 'shape': unsq_dims.shape},
'unsq_dims_const': {'value': unsq_dims, 'shape': unsq_dims.shape},
})
graph_ref = build_graph(self.nodes_attributes,
[('data_1', 'unsq'),
('unsq_dims_const', 'unsq_dims'),
('unsq_dims', 'unsq'),
('unsq', 'data_2')],
{'data_1': {'shape': np.array([2, 3, 64, 64])},
'unsq_dims': {'value': int64_array([4]), 'shape': unsq_dims.shape},
'unsq_dims_const': {'value': int64_array([4]), 'shape': unsq_dims.shape},
'data_2': {'shape': np.array([2, 3, 64, 64, 1])},
})
unsqueeze_node = Node(graph, 'unsq')
Unsqueeze.infer(unsqueeze_node)
(flag, resp) = compare_graphs(graph, graph_ref, 'data_2')
self.assertTrue(flag, resp)
def replace_op(self, graph: nx.MultiDiGraph, node: Node):
# reshape tensor with batch indices to 2d
unsqueeze_op = Unsqueeze(
graph, {'unsqueeze_dims': np.array([1], dtype=np.int64)})
unsqueeze_node = unsqueeze_op.create_node([node.in_node(2)])
concat_op = Concat(
graph, {
'axis': 1,
'name': node.name + '/concat_batch_indices_and_boxes'
})
concat_node = concat_op.create_node([unsqueeze_node, node.in_node(1)])
# do not remove edge with crop_size because it is needed in the partial infer
graph.remove_edge(node.in_node(1).id, node.id)
# input to the CropAndResize contains boxes coordinates in YXYX layout. But IE layer ROIPooling expects
# coordinates in the XYXY layout, so convolution is added here to swap coordinates
swapped_box_coordinates_node = add_convolution_to_swap_xy_coordinates(
graph, concat_node, 5)
# reshape locations tensor to 2D so it could be passed to Eltwise which will be converted to ScaleShift
reshape_2d_op = Reshape(graph, dict(dim=np.array([-1, 5])))
reshape_2d_node = reshape_2d_op.create_node(
[swapped_box_coordinates_node], dict(name='reshape_2d_'))
create_edge(reshape_2d_node, node, 0, 1)
# do not replace any output edge
return []
def add_unsqueeze_for_new(graph: Graph, ss_node: Node):
log.info(
"StridedSlice op with new axis mask '{}' has been detected".format(
ss_node.id))
if len(ss_node.in_nodes()) != 4 or len(ss_node.out_nodes()) != 1:
return
shape_out = ss_node.out_node().shape
dim = np.array(range(len(ss_node['new_axis_mask'])))[np.array(
ss_node['new_axis_mask'], dtype=bool)]
ss_shape = []
for i in range(0, len(ss_node['new_axis_mask'])):
if not ss_node['new_axis_mask'][i]:
ss_shape.append(shape_out[i])
else:
ss_node['new_axis_mask'][i] = 0
ss_node.out_port(0).data.set_shape(ss_shape)
# insert Unsqueeze
unsqueeze_node = Unsqueeze(graph,
dict(name=ss_node.name +
'/Unsqueeze_new')).create_node()
ss_node.out_port(0).get_connection().insert_node(unsqueeze_node)
unsqueeze_node.out_port(0).data.set_shape(shape_out)
dims_node = Const(graph, {
'name': unsqueeze_node.id + '/Indices',
'value': int64_array(dim)
}).create_node()
dims_node.out_port(0).connect(unsqueeze_node.in_port(1))
def test_unsqueeze_infer(self):
graph = build_graph(
self.nodes_attributes, [('data_1', 'unsq'), ('unsq', 'data_2')], {
'data_1': {
'shape': np.array([1, 3, 64, 64])
},
'unsq': {
'unsqueeze_dims': np.array([0, 4])
}
})
graph_ref = build_graph(
self.nodes_attributes, [('data_1', 'unsq'), ('unsq', 'data_2')], {
'data_1': {
'shape': np.array([1, 3, 64, 64])
},
'unsq': {
'unsqueeze_dims': np.array([0, 4])
},
'data_2': {
'shape': np.array([1, 1, 3, 64, 1, 64])
}
})
unsqueeze_node = Node(graph, 'unsq')
Unsqueeze.infer(unsqueeze_node)
(flag, resp) = compare_graphs(graph, graph_ref, 'data_2')
self.assertTrue(flag, resp)
def test_unsqueeze_infer(self, input_shape, unsq_dims, output_shape,
ref_uns_dims, input_value, output_value):
graph = build_graph(
self.nodes_attributes, [('data_1', 'unsq'),
('unsq_dims_const', 'unsq_dims'),
('unsq_dims', 'unsq'), ('unsq', 'data_2')],
{
'data_1': {
'shape': input_shape,
'value': input_value
},
'unsq_dims': {
'value': unsq_dims,
'shape': unsq_dims.shape
},
'unsq_dims_const': {
'value': unsq_dims,
'shape': unsq_dims.shape
},
})
graph_ref = build_graph(
self.nodes_attributes, [('data_1', 'unsq'),
('unsq_dims_const', 'unsq_dims'),
('unsq_dims', 'unsq'), ('unsq', 'data_2')],
{
'data_1': {
'shape': input_shape,
'value': input_value
},
'unsq_dims': {
'value': ref_uns_dims,
'shape': ref_uns_dims.shape
},
'unsq_dims_const': {
'value': ref_uns_dims,
'shape': ref_uns_dims.shape
},
'data_2': {
'shape': output_shape,
'value': output_value
},
})
unsqueeze_node = Node(graph, 'unsq')
Unsqueeze.infer(unsqueeze_node)
(flag, resp) = compare_graphs(graph, graph_ref, 'data_2')
self.assertTrue(flag, resp)
self.assertTrue(
strict_compare_tensors(
Node(graph, 'data_2').shape,
Node(graph_ref, 'data_2').shape))
if Node(graph_ref, 'data_2').value is not None:
self.assertTrue(
strict_compare_tensors(
Node(graph, 'data_2').value,
Node(graph_ref, 'data_2').value))
def extract(node):
axis = np.array(onnx_attr(node, 'axes', 'ints', default=[]),
dtype=np.int64)
attrs = {'unsqueeze_dims': axis if len(axis) != 0 else None}
# update the attributes of the node
Unsqueeze.update_node_stat(node, attrs)
return __class__.enabled
def replace_pattern(graph: Graph, match: dict):
fq = match['fq']
if len(fq.out_port(0).get_destinations()) > 1:
# FQ should have only one child -- Transpose for optimization
return
transpose = match['transpose']
name = fq.soft_get('name', fq.id)
input_shape = transpose.in_port(0).data.get_shape()
# detaching transpose from the graph
transpose.out_port(0).get_connection().set_source(transpose.in_port(0).get_connection().get_source())
transpose.in_port(0).disconnect()
for idx, port in fq.in_ports().items():
transpose_copy = transpose.copy_node({'override_output_shape': True})
transpose.in_port(1).get_source().connect(transpose_copy.in_port(1))
start_port = transpose_copy.in_port(0)
idxs = np.arange(len(input_shape) - len(port.data.get_shape()))
if idxs.size != 0:
axis = Const(graph, {'name': name + '/in_{}_unsqueeze_axis'.format(idx),
'value': int64_array(idxs)}).create_node()
unsqueeze = Unsqueeze(graph, {'name': name + '/in_{}_unsqueeze'.format(idx)}).create_node()
axis.out_port(0).connect(unsqueeze.in_port(1))
unsqueeze.out_port(0).connect(transpose_copy.in_port(0))
start_port = unsqueeze.in_port(0)
src = port.get_source()
port.get_connection().set_source(transpose_copy.out_port(0))
src.connect(start_port)
def replace_op(self, graph: Graph, node: Node):
if node.has_and_set('inputs_preprocessed'):
log.debug('Node "{}" has already been preprocessed'.format(
node.soft_get('name')))
return []
# reshape tensor with batch indices to 2d
unsqueeze_op = Unsqueeze(
graph, {'unsqueeze_dims': np.array([1], dtype=np.int64)})
unsqueeze_node = unsqueeze_op.create_node([node.in_node(2)])
concat_op = Concat(
graph, {
'axis': 1,
'name': node.name + '/concat_batch_indices_and_boxes',
'in_ports_count': 2
})
concat_node = concat_op.create_node([unsqueeze_node, node.in_node(1)])
# do not remove edge with crop_size because it is needed in the partial infer
graph.remove_edge(node.in_node(1).id, node.id)
# input to the CropAndResize contains boxes coordinates in YXYX layout. But IE layer ROIPooling expects
# coordinates in the XYXY layout, so convolution is added here to swap coordinates
swapped_box_coordinates_node = add_convolution_to_swap_xy_coordinates(
graph, concat_node, 5)
# reshape locations tensor to 2D so it could be passed to Eltwise which will be converted to ScaleShift
reshape_2d_op = Reshape(graph, dict(dim=np.array([-1, 5])))
reshape_2d_node = reshape_2d_op.create_node(
[swapped_box_coordinates_node],
dict(name=swapped_box_coordinates_node.id + '/reshape_2d_',
nchw_layout=True))
graph.create_edge(reshape_2d_node, node, 0, 1)
# do not replace any output edge
return []
def find_and_replace_pattern(self, graph: Graph):
for expand_dims_node in graph.get_op_nodes(op='ExpandDims'):
if len(expand_dims_node.in_nodes()) == 1:
expand_axis = expand_dims_node.expand_axis
if not isinstance(expand_axis, np.ndarray):
expand_axis = int64_array([expand_axis]).flatten()
unsqueeze_node = Unsqueeze(graph, {'name': expand_dims_node.id}).create_node()
unsqueeze_dims_node = Const(graph, {'name': expand_dims_node.id + '/Dims',
'value': expand_axis}).create_node()
expand_dims_node.in_port(0).get_connection().set_destination(unsqueeze_node.in_port(0))
expand_dims_node.out_port(0).get_connection().set_source(unsqueeze_node.out_port(0))
unsqueeze_node.in_port(1).connect(unsqueeze_dims_node.out_port(0))
else:
log.error('The ExpandDims node {} has more than 1 input'.format(expand_dims_node.soft_get('name')))
def replace_pattern(graph: Graph, match: dict):
"""
Workarounds not supported type of Tile in Inference Engine (Tiles are supported for 2-D or 4-D tensors):
Searches for Tiles with 3D shapes and covers it with Reshapes.
Example: Tile (axis=1, tiles=16):
in_shape: [1,1,101]
out_shape: [1,16,101]
Old behaviour:
Tile -> [1,16,101]
New behaviour:
Reshape [1,1,101,1] -> Tile -> [1,16,101,1] -> Reshape [1,16,101]
"""
node = match['tile']
name = node.soft_get('name', node.id)
out_shape = node.out_port(0).data.get_shape()
assert out_shape is not None, 'Output shape is undefined for {} in back phase'.format(name)
if out_shape.size != 3:
return
inp_shape = node.in_port(0).data.get_shape()
assert inp_shape is not None, 'Input shape is undefined for {} in back phase'.format(name)
unsqueeze_dim = Const(graph, {'name': name + '/3D_Tile_Unsqueeze_dim', 'value': int64_array([3])}).create_node()
unsqueeze = Unsqueeze(graph, {'name': name + '/3D_Tile_Unsqueeze', 'override_output_shape': True}).create_node()
unsqueeze_dim.out_port(0).connect(unsqueeze.in_port(1))
const = Const(graph, {'name': name + '/additional_axis', 'value': int64_array([1])}).create_node()
new_tiles = new_shape_node_from_shape_nodes([node.in_port(1).get_source().node, const])
node.in_port(1).get_connection().set_source(new_tiles.out_port(0))
squeeze_dim = Const(graph, {'name': name + '/3D_Tile_Squeeze_dim', 'value': int64_array([3])}).create_node()
squeeze = Squeeze(graph, {'name': name + '/3D_Tile_Squeeze', 'override_output_shape': True}).create_node()
squeeze_dim.out_port(0).connect(squeeze.in_port(1))
source = node.in_port(0).get_source()
node.in_port(0).get_connection().set_source(unsqueeze.out_port(0))
unsqueeze.in_port(0).connect(source)
node.out_port(0).get_connection().set_source(squeeze.out_port(0))
node.out_port(0).connect(squeeze.in_port(0))
node['override_output_shape'] = True
new_tiles['override_output_shape'] = True
node['need_shape_inference'] = True
def replace_pattern(graph, match: dict):
# Here we will found all parts of TI: condition, inputs/outputs, back edges, body and create TensorIterator Op
# and make all checks needed for TensorIterator work
cond_data = match['condition'].out_node(
0) if not match['condition'].out_port(0).disconnected() else None
time_data = match['condition'].out_node(1) if len(
match['condition'].out_nodes()) >= 1 else None
name = match['condition'].name
back_edges = []
inputs = []
outputs = []
if cond_data is not None:
for node in cond_data.out_nodes():
if node['kind'] == 'op' and node[
'op'] == 'TensorIteratorBackEdge':
back_edges.append(node.id)
elif node['kind'] == 'op' and node[
'op'] == 'TensorIteratorInput':
inputs.append(node.id)
elif node['kind'] == 'op' and node[
'op'] == 'TensorIteratorOutput':
outputs.append(node.id)
if time_data is not None:
for node in time_data.out_nodes():
if node['kind'] == 'op' and node['op'] == 'TensorIteratorInput':
inputs.append(node.id)
elif node['kind'] == 'op' and node[
'op'] == 'TensorIteratorOutput':
outputs.append(node.id)
else:
# something goes wrong here
assert False
condition = match['condition']
tensor_sequence_length = condition.in_node(0)
nodes_to_remove = [
n.id
for n in (condition, cond_data, time_data, tensor_sequence_length)
if n is not None
]
graph.remove_nodes_from(nodes_to_remove)
body_nodes, extra_inputs = get_body(graph, inputs, outputs)
if cond_data is not None:
body_nodes = list(set(body_nodes) - set([cond_data]))
inputs += extra_inputs
assert all([node in graph.nodes() for node in body_nodes])
inputs = [Node(graph, node) for node in inputs]
outputs = [Node(graph, node) for node in outputs]
back_edges = [Node(graph, node) for node in back_edges]
external_inputs = [{
'external_data_id':
node.in_node(1 if node.has_valid('axis') else 0),
'internal_data_id':
node.out_node(0),
'axis':
node.axis,
'start':
node.start,
'end':
node.end,
'stride':
node.stride,
'part_size':
node.part_size
} for node in inputs]
external_outputs = [{
'external_data_id':
node.out_node(0),
'internal_data_id':
node.in_node(1 if node.has_valid('axis') else 0),
'axis':
node.axis,
'start':
node.start,
'end':
node.end,
'stride':
node.stride,
'part_size':
node.part_size
} for node in outputs]
back_edges_data = [{
'from_data_id': node.in_node(1),
'to_data_id': node.out_node(0),
'init_data_id': node.in_node(0),
} for node in back_edges]
body = Graph(name='body')
body.graph = graph.graph
body.add_nodes_from([(node, graph.node[node]) for node in body_nodes])
body.add_edges_from([
(u, v, k, d) for u, v, k, d in graph.edges(data=True, keys=True)
if u in body_nodes and v in body_nodes
])
graph.remove_nodes_from(body_nodes + [match['condition'].id] +
[inp.id for inp in inputs] +
[out.id for out in outputs])
internal_id_count = 0
real_back_edges = []
for edge in back_edges_data:
assert edge['from_data_id'].id in body.nodes()
assert edge['to_data_id'].id in body.nodes()
assert edge['init_data_id'].id in body.nodes()
edge['from_data_id'] = Node(body, edge['from_data_id'].id)
edge['to_data_id'] = Node(body, edge['to_data_id'].id)
edge['init_data_id'] = Node(body, edge['init_data_id'].id)
add_opoutput(body, edge['from_data_id'].id, 0, False)
# Assign/reuse ids for the back-edge start; it comes from from_data_id
assert len(edge['from_data_id'].in_nodes()) == 1
# layer id
if not edge['from_data_id'].in_node().has_valid(
'internal_layer_id'):
edge['from_data_id'].in_node(
)['internal_layer_id'] = internal_id_count
internal_id_count += 1
edge['from_layer'] = edge['from_data_id'].in_node(
)['internal_layer_id']
# port id
if 'internal_port_id' not in edge['from_data_id'].in_edge():
edge['from_data_id'].in_edge(
)['internal_port_id'] = internal_id_count
internal_id_count += 1
edge['from_port'] = edge['from_data_id'].in_edge(
)['internal_port_id']
# Look at all consumers for a data that ends a back-edge
# For each such consumer, there will be a separate back-edge (and input)
current_real_back_edges = []
for _, consumer, key, edge_attrs in body.out_edges(
edge['to_data_id'].id, data=True, keys=True):
real_edge = {}
real_edge.update(
edge) # all real back_edges have the same back-edge start
consumer = Node(body, consumer)
if real_edge['to_data_id'].in_node().has_valid(
'internal_layer_id'):
assert False
real_edge['to_data_id'].out_node()['internal_layer_id'] = \
real_edge['to_data_id'].in_node().internal_layer_id
elif not consumer.has_valid('internal_layer_id'):
consumer['internal_layer_id'] = internal_id_count
internal_id_count += 1
real_edge['to_layer'] = consumer['internal_layer_id']
assert 'internal_port_id' not in edge_attrs
assert len(real_edge['init_data_id'].out_edges()) == 1
assert not 'internal_port_id' in real_edge[
'init_data_id'].out_edge()
edge_attrs['internal_port_id'] = internal_id_count
internal_id_count += 1
real_edge['to_port'] = edge_attrs['internal_port_id']
real_edge['consumer'] = consumer
real_edge['consumer_key'] = key
real_edge['attrs'] = deepcopy(edge_attrs)
current_real_back_edges.append(real_edge)
# connect initial data node with each consumer providing actual edge attributes
body.add_edges_from([
(real_edge['init_data_id'].id, real_edge['consumer'].id,
real_edge['consumer_key'], real_edge['attrs'])
for real_edge in current_real_back_edges
])
body.remove_nodes_from(
[edge['to_data_id'].id, edge['to_data_id'].in_node().id])
real_back_edges += current_real_back_edges
real_external_inputs = []
for ext_inp in external_inputs:
assert ext_inp['external_data_id'].id not in body.nodes()
assert ext_inp['internal_data_id'].id in body.nodes()
ext_inp['internal_data_id'] = Node(body,
ext_inp['internal_data_id'].id)
if ext_inp['axis'] is not None:
# Insert squeezing resize at input port that has partitioning
shape = ext_inp['internal_data_id'].shape.copy()
assert not ext_inp['internal_data_id'].has_valid('value')
new_input_data = Op._create_data_node(
body,
ext_inp['internal_data_id'].name + '/UnsqueezedInput',
dict(shape=shape_insert(shape, ext_inp['axis'], 1)))
reshape_op = Squeeze(
body,
dict(name=ext_inp['internal_data_id'].name +
'/InputSqueeze'))
reshape_dim_data = Const(
body, {
'name':
ext_inp['internal_data_id'].name + '/ReshapeDim',
'value': ext_inp['axis']
}).create_node_with_data()
reshape_op.create_node_with_data(
[new_input_data, reshape_dim_data],
data_nodes=[ext_inp['internal_data_id']])
ext_inp['internal_data_id'] = new_input_data
ext_inp['internal_data_id']['is_input'] = True
assert len(ext_inp['internal_data_id'].in_nodes()) == 0
ext_inp['external_port_id'] = internal_id_count
internal_id_count += 1
for _, consumer, edge_attrs in body.out_edges(
ext_inp['internal_data_id'].id, data=True):
real_ext_inp = {}
real_ext_inp.update(ext_inp)
consumer = Node(body, consumer)
if not consumer.has_valid('internal_layer_id'):
consumer['internal_layer_id'] = internal_id_count
internal_id_count += 1
if not 'internal_port_id' in edge_attrs:
edge_attrs['internal_port_id'] = internal_id_count
internal_id_count += 1
real_ext_inp['internal_layer_id'] = consumer[
'internal_layer_id']
real_ext_inp['internal_port_id'] = edge_attrs[
'internal_port_id']
real_external_inputs.append(real_ext_inp)
for ext_out in external_outputs:
assert ext_out['external_data_id'].id not in body.nodes()
assert ext_out['internal_data_id'].id in body.nodes()
ext_out['internal_data_id'] = Node(body,
ext_out['internal_data_id'].id)
if ext_out['axis'] is not None:
# Insert unsqueezing resize at output port that has partitioning
reshape_op = Unsqueeze(
body,
dict(name=ext_out['internal_data_id'].name +
'/OutputUnsqueeze'))
reshape_dim_data = Const(
body, {
'name':
ext_out['internal_data_id'].name + '/ReshapeDim',
'value': ext_out['axis']
}).create_node_with_data()
ext_out['internal_data_id'] = reshape_op.create_node_with_data(
[ext_out['internal_data_id'], reshape_dim_data])
# TODO: add here working with simple outputs
if not any([
out_node.soft_get('op', None) == 'Result'
for out_node in ext_out['internal_data_id'].out_nodes()
]):
add_opoutput(body, ext_out['internal_data_id'].id, 0, False)
# assert len(ext_out['internal_data_id'].out_nodes()) == 0
assert len(ext_out['internal_data_id'].in_nodes()) == 1
if not 'internal_layer_id' in ext_out['internal_data_id'].in_node(
):
ext_out['internal_data_id'].in_node(
)['internal_layer_id'] = internal_id_count
internal_id_count += 1
if not 'internal_port_id' in ext_out['internal_data_id'].in_edge():
ext_out['internal_data_id'].in_edge(
)['internal_port_id'] = internal_id_count
internal_id_count += 1
ext_out['internal_layer_id'] = ext_out['internal_data_id'].in_node(
)['internal_layer_id']
ext_out['internal_port_id'] = ext_out['internal_data_id'].in_edge(
)['internal_port_id']
ext_out['external_port_id'] = internal_id_count
internal_id_count += 1
# create TensorIterator layer with pre-computed components
ti_op = TensorIterator(
graph, {
'name':
name + '/TensorIterator',
'body':
body,
'in_ports_count':
len(external_inputs),
'out_ports_count':
len(external_outputs),
'input_port_map': [{
field: external_input[field]
for field in [
'external_port_id', 'internal_layer_id',
'internal_port_id', 'axis', 'stride', 'part_size',
'start', 'end'
]
} for external_input in real_external_inputs],
'output_port_map': [{
field: external_output[field]
for field in [
'external_port_id', 'internal_layer_id',
'internal_port_id', 'axis', 'stride', 'part_size',
'start', 'end'
]
} for external_output in external_outputs],
'back_edges': [{
field: edge[field]
for field in
['from_layer', 'from_port', 'to_layer', 'to_port']
} for edge in real_back_edges],
})
ti_outs = ti_op.create_node_with_data(
inputs=[inp['external_data_id'] for inp in external_inputs],
edge_attrs=[{
'external_port_id': inp['external_port_id']
} for inp in external_inputs],
data_nodes=[out['external_data_id'] for out in external_outputs])
if not isinstance(ti_outs, list):
ti_outs = [ti_outs]
for i, out in enumerate(ti_outs):
out.in_edge(
)['external_port_id'] = external_outputs[i]['external_port_id']
ti = ti_outs[0].in_node()
TensorIterator.cover_body_input_data_nodes_with_parameter_ops(ti)
TensorIterator.cover_body_constant_data_nodes_with_const_ops(ti)
TensorIterator.normalize_internal_ids(ti)
def mxrepeat_decomposition(node: Node):
graph = node.graph
name = node.soft_get('name', node.id)
rename_node(node, name + '/to_be_removed')
# Unqueeze
input_rank = Rank(graph, {'name': name + '/Rank'}).create_node()
node.in_port(0).get_source().connect(input_rank.in_port(0))
axis = get_canonical_axis_index_node(input_rank, node.axis)
unsqueeze_axis = create_op_node_with_second_input(
graph,
Add,
int64_array([1]), {'name': name + '/Unsqueeze/Axis'},
input_node=axis)
unsqueeze = Unsqueeze(graph, {
'name': name + '/Unsqueeze'
}).create_node()
unsqueeze.in_port(1).connect(unsqueeze_axis.out_port(0))
# Tile (1, 1, ..., repeats, ..., 1)
# we generate tile array according to the following table:
# parts: | first | repeats | second |
# i: | 0, 1, ..., axis,| axis + 1,| ..., rank+1 |
# tile_array: | 1, 1, ..., 1 ,| repeats ,| ..., 1 |
one = Const(graph, {
'name': name + '/Broadcast/One',
'value': int64_array([1])
}).create_node()
first_ones = Broadcast(graph, {
'name': name + '/Broadcast/Ones_first_part'
}).create_node()
first_ones.in_port(0).connect(one.out_port(0))
first_ones.in_port(1).connect(unsqueeze_axis.out_port(0))
repeats = Const(graph, {
'name': name + '/repeats',
'value': int64_array([node.repeats])
}).create_node()
second_ones = Broadcast(graph, {
'name': name + '/Broadcast/Ones_second_part'
}).create_node()
second_part_broadcast_shape = Sub(
graph, {
'name': name + '/Broadcast/Shape/second_part'
}).create_node()
second_part_broadcast_shape.in_port(0).connect(input_rank.out_port(0))
second_part_broadcast_shape.in_port(1).connect(
unsqueeze_axis.out_port(0))
second_ones.in_port(0).connect(one.out_port(0))
second_ones.in_port(1).connect(second_part_broadcast_shape.out_port(0))
tile_repeats = new_shape_node_from_shape_nodes(
[first_ones, repeats, second_ones])
tile = Tile(graph, {'name': name + '/Tile'}).create_node()
tile.in_port(1).connect(tile_repeats.out_port(0))
# Reshape (input_shape[:axis], input_shape[axis] * repeats, input_shape[axis+1:])
# we generate reshape dim array according to the following table:
# parts: | first | rep | second |
# i: | 0, 1, ... ,| axis, | ..., rank |
# dim_array: | inp_sh[i] ,| input_shape[axis] * repeats ,| inp_sh[i] |
input_shape = Shape(graph, {'name': name + '/Shape'}).create_node()
node.in_port(0).get_source().connect(input_shape.in_port(0))
first_input_shape_part = get_shape_values_by_range_idxs(
input_shape,
input_rank,
begin=0,
end=node.axis,
include_begin=True,
include_end=False)
original_axis_dim = create_op_with_const_inputs(
graph,
Gather, {2: int64_array(0)}, {'name': name + '/OriginalDim'},
input_node=input_shape)
original_axis_dim.in_port(1).connect(axis.out_port(0))
repeated_dimention = Mul(graph, {
'name': name + '/RepeatedDim'
}).create_node()
repeated_dimention.in_port(0).connect(original_axis_dim.out_port(0))
repeated_dimention.in_port(1).connect(repeats.out_port(0))
second_input_shape_part = get_shape_values_by_range_idxs(
input_shape,
input_rank,
begin=node.axis,
end=-1,
include_begin=False,
include_end=True)
output_shape = new_shape_node_from_shape_nodes([
first_input_shape_part, repeated_dimention, second_input_shape_part
])
reshape = Reshape(graph, {'name': name}).create_node()
rename_node(reshape, name)
reshape.in_port(1).connect(output_shape.out_port(0))
# Final connections
node.in_port(0).get_connection().set_destination(unsqueeze.in_port(0))
tile.in_port(0).connect(unsqueeze.out_port(0))
reshape.in_port(0).connect(tile.out_port(0))
node.out_port(0).get_connection().set_source(reshape.out_port(0))
def replace_pattern(self, graph: Graph, match: dict):
lstm = match['lstm']
# Build TensorIterator body first
body = Graph(name=lstm.name + '/sub_graph')
body.graph = graph.graph
# 1. Input squeeze Reshape
inputs = [
Op._create_data_node(
body, lstm.name + '/inport/' + str(inp), {
'shape':
lstm.in_node(inp).shape.copy(),
'value':
lstm.in_node(inp).value.copy() if lstm.in_node(inp).value
is not None and inp in [1, 2] else None
}) for inp in [0, 4, 5, 1, 2]
] # X, WR, B, h_init, c_init
inputs[0].shape[lstm.sequence_dim] = 1
input_squeeze = Squeeze(
body, dict(name=lstm.name + '/input_squeeze', internal_layer_id=0))
squeeze_dim_data = Const(body, {
'name': lstm.name + '/input_squeeze_dim',
'value': [lstm.sequence_dim]
}).create_node_with_data()
inputs[0] = input_squeeze.create_node_with_data(
[inputs[0], squeeze_dim_data],
edge_attrs=[{
'internal_port_id': 0
}])
# 2. Output unsqueeze Reshape
outputs = [
Op._create_data_node(
body, lstm.name + '/outport/' + str(out), {
'shape':
lstm.out_node(out).shape.copy() if out in lstm.out_nodes()
else lstm.in_node(4).shape.copy()
}) for out in [0, 1]
]
for out in outputs:
add_opoutput(body, out.id, 0, False)
outputs[0].shape = shape_delete(outputs[0].shape, lstm.sequence_dim)
output_unsqueeze = Unsqueeze(
body, dict(name=lstm.name + 'output_unsqueeze',
internal_layer_id=2))
unsqueeze_dim_data = Const(
body, {
'name': lstm.name + '/output_unsqueeze_dim',
'value': [lstm.sequence_dim]
}).create_node_with_data()
# 3. LSTMCell
lstm_cell_op = LSTMCell(
body,
dict(hidden_size=lstm.hidden_size,
activations=lstm.activations,
activation_alpha=lstm.activation_alpha,
activation_beta=lstm.activation_beta,
clip=lstm.clip,
input_forget=lstm.input_forget,
name=lstm.name + '/LSTMCell',
internal_layer_id=1))
lstm_cell_node = lstm_cell_op.create_node_with_data(
inputs,
data_nodes=outputs,
edge_attrs=[{}, {
'internal_port_id': 1
}, {
'internal_port_id': 2
}, {
'bin': 'weights'
}, {
'bin': 'biases'
}])
lstm_cell_node[0].in_node().out_edge(0)['internal_port_id'] = 4
lstm_cell_node[0].in_node().out_edge(1)['internal_port_id'] = 5
lstm_cell_node[0] = output_unsqueeze.create_node_with_data(
[lstm_cell_node[0], unsqueeze_dim_data])
lstm_cell_node[0].in_node().out_edge(0)['internal_port_id'] = 3
add_opoutput(body, lstm_cell_node[0].id, 0, False)
# 4. TensorIterator layer creating
assert lstm.direction in ['forward', 'reverse']
if lstm.direction == 'forward':
stride = 1
start = None
end = None
else:
assert lstm.direction == 'reverse'
stride = -1
start = -1
end = 0
output_port_map = [{
'external_port_id': 3,
'internal_layer_id': 2,
'internal_port_id': 3,
'axis': lstm.sequence_dim,
'stride': stride,
'start': start,
'end': end,
'part_size': 1,
}]
# Adding h_state, c_state to outputs
if len(lstm.out_nodes()) == 3:
output_port_map.extend([{
'external_port_id': 4,
'internal_layer_id': 1,
'internal_port_id': 4,
}, {
'external_port_id': 5,
'internal_layer_id': 1,
'internal_port_id': 5,
}])
ti_op = TensorIterator(
graph, {
'name':
lstm.name + '/TensorIterator',
'body':
body,
'in_ports_count':
3,
'out_ports_count':
len(lstm.out_nodes()),
'input_port_map': [
{
'external_port_id': 0,
'internal_layer_id': 0,
'internal_port_id': 0,
'axis': lstm.sequence_dim,
'stride': stride,
'start': start,
'end': end,
'part_size': 1,
},
{
'external_port_id': 1,
'internal_layer_id': 1,
'internal_port_id': 1,
},
{
'external_port_id': 2,
'internal_layer_id': 1,
'internal_port_id': 2,
},
],
'output_port_map':
output_port_map,
'back_edges': [
{
'from_layer': 1,
'from_port': 4,
'to_layer': 1,
'to_port': 1,
},
{
'from_layer': 1,
'from_port': 5,
'to_layer': 1,
'to_port': 2,
},
]
})
assert sorted(lstm.out_nodes().keys()) == list(range(len(lstm.out_nodes()))), \
"There are gaps in output ports of LSTMSequence operation. Node {}".format(lstm.id)
outs = ti_op.create_node_with_data(
[lstm.in_node(i) for i in [0, 4, 5]], # X, h_init, c_init
data_nodes=[
lstm.out_node(i) for i in range(len(lstm.out_nodes()))
],
edge_attrs=[{
'external_port_id': 0
}, {
'external_port_id': 1
}, {
'external_port_id': 2
}])
if not isinstance(outs, list):
outs = list([outs])
graph.remove_node(lstm.id)
outs[0].in_edge(0)['external_port_id'] = 3
for i, out in enumerate(outs[1:]):
external_port_id = 4 + i
out.in_edge()['external_port_id'] = external_port_id
ti = outs[0].in_node()
TensorIterator.cover_body_input_data_nodes_with_parameter_ops(ti)
TensorIterator.cover_body_constant_data_nodes_with_const_ops(ti)
TensorIterator.normalize_internal_ids(ti)
def replace_pattern(graph: Graph, match: dict):
node = match['matmul']
name = node.soft_get('name', node.id)
A_shape = node.in_port(0).data.get_shape()
B_shape = node.in_port(1).data.get_shape()
out_shape = node.out_port(0).data.get_shape()
assert A_shape is not None and B_shape is not None and out_shape is not None
B_value = node.in_port(1).data.get_value()
if (B_value is not None or node.in_port(1).get_source().node.has_and_set('stop_value_propagation')) and B_shape[
B_shape != 1].size <= 2:
# transferring from MatMul representation: [B, I, K] * [B, K, O] = [B, I, O]
# to FullyConnected representation: [I, K] * [O, K] = [I, O]
B, I, K, O, aligned_A_shape, aligned_B_shape = MatMulToFullyConnected.get_matmul_BIKO(node)
# weights normalization
if not node.transpose_b:
# FullyConnected weights layout is OI
# MatMul second input layout is (B)IO
transpose_order = list(range(B_shape.size))
transpose_order[-1], transpose_order[-2] = transpose_order[-2], transpose_order[-1]
order = Const(graph, {'value': int64_array(transpose_order)}).create_node()
transpose = Transpose(graph, {'name': name + '/weights_transpose'}).create_node()
weights_source = node.in_port(1).get_source()
node.in_port(1).get_connection().set_source(transpose.out_port(0))
transpose.in_port(0).connect(weights_source)
transpose.in_port(1).connect(order.out_port(0))
order.infer(order)
transpose.infer(transpose)
if node.in_port(1).data.get_shape().size != 2:
const = Const(graph, {'value': int64_array([-1, K])}).create_node()
reshape = Reshape(graph, {'name': name + '/weights_reshape'}).create_node()
weights_source = node.in_port(1).get_source()
node.in_port(1).get_connection().set_source(reshape.out_port(0))
reshape.in_port(0).connect(weights_source)
reshape.in_port(1).connect(const.out_port(0))
const.infer(const)
reshape.infer(reshape)
assert np.all(np.array_equal(node.in_port(1).data.get_shape(), int64_array([O, K]))), \
"MatMul `{}` was not converted to FullyConnected: wrong weights shape: {}, " \
"B={}, I={}, K={}, O={}".format(name, node.in_port(1).data.get_shape(), B, I, K, O)
node.in_port(1).bin = 'weights'
del node['transpose_b']
# input normalization
if node.transpose_a:
transpose_order = list(range(A_shape.size))
transpose_order[-1], transpose_order[-2] = transpose_order[-2], transpose_order[-1]
order = Const(graph, {'value': int64_array(transpose_order)}).create_node()
transpose = Transpose(graph, {'name': name + '/input_transpose'}).create_node()
input_source = node.in_port(0).get_source()
node.in_port(0).get_connection().set_source(transpose.out_port(0))
transpose.in_port(0).connect(input_source)
transpose.in_port(1).connect(order.out_port(0))
order.infer(order)
transpose.infer(transpose)
if A_shape.size != 2:
const = Const(graph, {'value': int64_array([-1, K])}).create_node()
reshape = Reshape(graph, {'name': name + '/input_reshape'}).create_node()
input_source = node.in_port(0).get_source()
node.in_port(0).get_connection().set_source(reshape.out_port(0))
reshape.in_port(0).connect(input_source)
reshape.in_port(1).connect(const.out_port(0))
const.infer(const)
reshape.infer(reshape)
assert np.all(np.array_equal(node.in_port(0).data.get_shape(), int64_array([np.prod(B) * I, K]))), \
"MatMul `{}` wasn't converted to FullyConnected: wrong input shape: {}, " \
"B={}, I={}, K={}, O={}".format(name, node.in_port(0).data.get_shape(), B, I, K, O)
del node['transpose_a']
FullyConnected.update_node_stat(node, {'out-size': O})
# output normalization
if out_shape.size != 2:
const = Const(graph, {'value': int64_array([*B, I, O])}).create_node()
reshape = Reshape(graph, {'name': name + '/output_reshape'}).create_node()
dst = node.out_port(0).get_destination()
node.out_port(0).get_connection().set_destination(reshape.in_port(0))
const.out_port(0).connect(reshape.in_port(1))
reshape.out_port(0).connect(dst)
node.infer(node)
const.infer(const)
reshape.infer(reshape)
else:
assert A_shape.size == out_shape.size
assert B_shape.size <= out_shape.size
if B_shape.size != out_shape.size:
unsqueeze_dim = Const(graph, {'value': int64_array(list(range(out_shape.size - B_shape.size)))
}).create_node()
unsqueeze = Unsqueeze(graph, {}).create_node()
B_source = node.in_port(1).get_source()
node.in_port(1).get_connection().set_source(unsqueeze.out_port(0))
unsqueeze.in_port(0).connect(B_source)
unsqueeze.in_port(1).connect(unsqueeze_dim.out_port(0))
unsqueeze_dim.infer(unsqueeze_dim)
unsqueeze.infer(unsqueeze)
Gemm.update_node_stat(node, {
'transpose_a': node.has_and_set('transpose_a'),
'transpose_b': node.has_and_set('transpose_b'),
})
def replace_pattern(self, graph: Graph, match: dict):
if match['rnn_layer']['op'] == 'LSTM':
return
rnn_layer = match['rnn_layer']
# Build TensorIterator body first
body = Graph(name=rnn_layer.name + '/sub_graph')
body.graph = graph.graph
# 1. Input squeeze Reshape
inputs = [
Op._create_data_node(
body, rnn_layer.name + '/inport/' + str(inp), {
'shape':
rnn_layer.in_node(inp).shape.copy(),
'value':
rnn_layer.in_node(inp).value.copy()
if rnn_layer.in_node(inp).value is not None
and inp in [1, 2] else None
}) for inp in [0, 4, 1, 2]
] # X, h_init, WR, B
inputs[0].shape[rnn_layer.sequence_dim] = 1
input_squeeze = Squeeze(
body,
dict(name=rnn_layer.name + '/input_squeeze', internal_layer_id=0))
input_squeeze_dim = Const(
body,
dict(name=rnn_layer.name + '/input_squeeze_dim',
value=rnn_layer.sequence_dim)).create_node_with_data()
inputs[0] = input_squeeze.create_node_with_data(
[inputs[0], input_squeeze_dim],
edge_attrs=[{
'internal_port_id': 0
}])
# 2. Output unsqueeze Reshape
outputs = [
Op._create_data_node(
body, rnn_layer.name + '/outport/' + str(out), {
'shape':
rnn_layer.out_node(out).shape.copy()
if out in rnn_layer.out_nodes() else None
}) for out in [0]
]
for out in outputs:
add_opoutput(body, out.id, 0, False)
outputs[0].shape = np.delete(outputs[0].shape.copy(),
rnn_layer.sequence_dim)
output_unsqueeze_dim = Const(
body,
dict(name=rnn_layer.name + '/output_unsqueeze_dim',
value=rnn_layer.sequence_dim)).create_node_with_data()
output_unsqueeze = Unsqueeze(
body,
dict(name=rnn_layer.name + '/output_unsqueeze/',
internal_layer_id=2))
additional_attrs = dict(activations=rnn_layer.activations,
activation_alpha=rnn_layer.activation_alpha,
activation_beta=rnn_layer.activation_beta,
clip=rnn_layer.clip)
if rnn_layer.op == 'GRU':
additional_attrs[
'linear_before_reset'] = rnn_layer.linear_before_reset
# 3. ***Cell
rnn_cell_op = self.get_rnn_cell(rnn_layer['op'])(
body,
dict(hidden_size=rnn_layer.hidden_size,
name=rnn_layer.name + '/{}Cell'.format(rnn_layer.op),
**additional_attrs,
internal_layer_id=1))
gru_cell = rnn_cell_op.create_node_with_data(inputs,
data_nodes=outputs,
edge_attrs=[{}, {
'internal_port_id':
1
}, {
'internal_port_id':
2
}, {
'bin':
'weights'
}, {
'bin':
'biases'
}])
# internal ports for outputs of cell
gru_cell.in_node().out_edge(0)['internal_port_id'] = 4 # h_state
gru_cell = output_unsqueeze.create_node_with_data(
[gru_cell, output_unsqueeze_dim])
gru_cell.in_node().out_edge(0)['internal_port_id'] = 3
add_opoutput(body, gru_cell.id, 0, False)
# 4. TensorIterator layer creating
assert rnn_layer.direction in ['forward', 'reverse']
if rnn_layer.direction == 'forward':
stride = 1
start = None
end = None
else:
assert rnn_layer.direction == 'reverse'
stride = -1
start = -1
end = 0
# stacked h_state
output_port_map = [{
'external_port_id': 3,
'internal_layer_id': 2,
'internal_port_id': 3,
'axis': rnn_layer.sequence_dim,
'stride': stride,
'start': start,
'end': end,
'part_size': 1,
}]
# Adding last h_state to outputs
if len(rnn_layer.out_nodes()) == 2:
output_port_map.extend([{
'external_port_id': 4,
'internal_layer_id': 1,
'internal_port_id': 4,
}])
ti_op = TensorIterator(
graph,
{
'name':
rnn_layer.name + '/TensorIterator',
'body':
body,
'in_ports_count':
4,
'out_ports_count':
len(rnn_layer.out_nodes()),
'input_port_map': [
{
'external_port_id': 0,
'internal_layer_id': 0,
'internal_port_id': 0,
'axis': rnn_layer.sequence_dim,
'stride': stride,
'start': start,
'end': end,
'part_size': 1,
},
{
'external_port_id': 1,
'internal_layer_id': 1,
'internal_port_id': 1,
},
],
'output_port_map':
output_port_map,
# only for h state
'back_edges': [
{
'from_layer': 1,
'from_port': 4,
'to_layer': 1,
'to_port': 1,
},
]
})
assert sorted(rnn_layer.out_nodes().keys()) == list(range(len(rnn_layer.out_nodes()))), \
"There are gaps in output ports of GRUSequence operation. Node {}".format(rnn_layer.id)
outs = ti_op.create_node_with_data(
[rnn_layer.in_node(i) for i in [0, 4]], # X, h_init
data_nodes=[
rnn_layer.out_node(i)
for i in range(len(rnn_layer.out_nodes()))
],
edge_attrs=[{
'external_port_id': 0
}, {
'external_port_id': 1
}])
if not isinstance(outs, list):
outs = list([outs])
graph.remove_node(rnn_layer.id)
outs[0].in_edge(0)['external_port_id'] = 3
for i, out in enumerate(outs[1:]):
external_port_id = 4 + i
out.in_edge()['external_port_id'] = external_port_id
ti = outs[0].in_node()
TensorIterator.cover_body_input_data_nodes_with_parameter_ops(ti)
TensorIterator.cover_body_constant_data_nodes_with_const_ops(ti)
TensorIterator.normalize_internal_ids(ti)
def extract(cls, node: Node):
Unsqueeze.update_node_stat(node, {})
return cls.enabled
def extract(node: Node):
Unsqueeze.update_node_stat(node, {})
return __class__.enabled