def check_init_states(graph: Graph, match: dict):
"""
Check if cell have initial states and create zeros states if not.
"""
rnn_layer = match['rnn_layer']
num_directions = 2 if rnn_layer.direction == 'bidirectional' else 1
batch_size = rnn_layer.in_node(0).shape[rnn_layer.batch_dim]
h_init_port = 5
c_init_port = 6
if h_init_port not in rnn_layer.in_nodes():
h_shape = [num_directions, batch_size, rnn_layer.hidden_size] # from ONNX spec
h_init = np.full(h_shape, 0, dtype=np.float32)
Op.create_and_connect_input_data_node(
graph,
rnn_layer,
{'value': h_init, 'shape': int64_array(h_init.shape)},
{'in': h_init_port, 'permutation': None}
)
if rnn_layer.op == 'LSTM':
if c_init_port not in rnn_layer.in_nodes():
c_shape = [num_directions, batch_size, rnn_layer.hidden_size] # from ONNX spec
c_init = np.full(c_shape, 0, dtype=np.float32)
Op.create_and_connect_input_data_node(
graph,
rnn_layer,
{'value': c_init, 'shape': int64_array(c_init.shape)},
{'in': c_init_port, 'permutation': None}
)
def insert_reduce(self,
model_graph,
insert_op,
node,
granularity,
type_stat,
node_name,
axis=1):
axis_const = self.find_axis(node, granularity, axis)
if isinstance(axis_const, str):
return (True, node.name)
out_port = self.get_out_port(node_name)
if out_port is not None:
node_name = f'{node_name[0]}.{out_port}'
reduce_op = create_op_node_with_second_input(
node.graph, insert_op, int64_array(axis_const),
dict(name=f'{type_stat}_{node_name}'))
reduce_op['fullname'] = reset_node_fullname(node.fullname,
reduce_op.name)
if node.graph != model_graph:
Op.create_data_node(reduce_op.graph, reduce_op, {'shape': [1]})
node.out_port(out_port if out_port else 0).connect(
reduce_op.in_port(0))
return self.insert_result(model_graph, node, reduce_op, type_stat,
out_port)
def repack_weights(self, graph: Graph, match: dict):
# Concat W, R in IE- format
# Delete useless num_dir dimensions and n_cells dimensions in W, R, B (peepholes?)
lstm = match['rnn_layer']
W, R, B = match['W'].value.copy(), match['R'].value.copy(), match['B'].value.copy()
graph.remove_edge(match['W'].id, lstm.id)
graph.remove_edge(match['R'].id, lstm.id)
graph.remove_edge(match['B'].id, lstm.id)
# Sum component of B that correspond to W and R
if lstm.op == 'GRU' and lstm.linear_before_reset:
B_shape = np.array(B.shape)
B_shape[3] = 4
B_shape[2] = 1
B_tmp = np.zeros(shape=B_shape)
B_tmp[:, :, :, 0, :] = B[:, :, 0, 0, :] + B[:, :, 1, 0, :]
B_tmp[:, :, :, 1, :] = B[:, :, 0, 1, :] + B[:, :, 1, 1, :]
B_tmp[:, :, :, 2, :] = B[:, :, 0, 2, :][:, :, np.newaxis, :]
B_tmp[:, :, :, 3, :] = B[:, :, 1, 2, :][:, :, np.newaxis, :]
B = B_tmp
else:
B = np.sum(B, axis=2, keepdims=True)
# Concatenate W, R to IE-compatible format
assert len(W.shape) == 5
assert len(R.shape) == 5
WR = np.concatenate([W, R], axis=4)
# Squeeze useless dimensions
assert WR.shape[0] == 1 # num_dir == 1
assert WR.shape[1] == 1 # num_cells == 1
assert B.shape[0] == 1
assert B.shape[1] == 1
WR = WR.squeeze(axis=(0, 1))
B = B.squeeze(axis=(0, 1))
# Flatten all output (0, 1) and input dimensions (2, 3)
final_shape_WR = [WR.shape[0] * WR.shape[1], -1]
assert final_shape_WR[0] == lstm.hidden_size * lstm.multiplier
WR = WR.reshape(final_shape_WR)
final_shape_B = final_shape_WR
if lstm.op == 'GRU' and lstm.linear_before_reset:
final_shape_B[0] = lstm.hidden_size * 4
B = B.reshape(final_shape_B)
# Squeeze fake dimension in B
B = B.squeeze(axis=-1)
for blob, port, name in [(WR, 1, 'weights'), (B, 2, 'biases')]:
Op.create_and_connect_input_data_node(
graph,
lstm,
{'value': blob, 'shape': np.array(blob.shape, dtype=np.int64)},
{'in': port, 'bin': name, 'permutation': None}
)
def extract(cls, node):
attrs = {
'data_type': tf_dtype_extractor(node.pb.attr["dtype"].type),
'shape': tf_tensor_shape(node.pb.attr["shape"].shape),
'identity': True,
'infer': lambda node: copy_shape_infer(node, value_infer=copy_value),
}
Op.update_node_stat(node, attrs)
return cls.enabled
def _create_data_if_necessary(self):
if self.node.graph.stage == 'front':
raise Error("_create_data_if_necessary method is not applicable for front Graph phase!")
if self.type == 'in':
raise Error("_create_data_if_necessary method is not applicable for 'in' Port type!")
if self.idx not in self.node.out_nodes(control_flow=self.control_flow):
from openvino.tools.mo.ops.op import Op
Op.create_data_node(self.node.graph, self.node, out_port=self.idx)
self.node['need_shape_inference'] = True
return self.node.out_node(self.idx, control_flow=self.control_flow)
def repack_weights(graph: Graph, match: dict):
"""
Repack weights into general format (described above) and reorder gates.
"""
rnn_layer = match['rnn_layer']
W = match['W'].value.copy()
R = match['R'].value.copy()
num_directions = 2 if rnn_layer.direction == 'bidirectional' else 1
graph.remove_edge(match['W'].id, rnn_layer.id)
graph.remove_edge(match['R'].id, rnn_layer.id)
# find optional 'B' biases blob
if 3 in rnn_layer.in_nodes():
# TODO: check if 'bin': 'B' attribute is assigned to this edge
B = rnn_layer.in_node(3).value.copy()
graph.remove_edge(rnn_layer.in_node(3).id, rnn_layer.id)
else:
B_shape = [num_directions, 2 * rnn_layer.multiplier * rnn_layer.hidden_size] # from ONNX spec
B = np.full(B_shape, 0, dtype=np.float32)
# Add extra dimensions for W, R and B for easier repacking and reordering
B = B.reshape([
num_directions, # 0: num of directions
rnn_layer.num_layers, # 1: num_layers
2, # 2: two input parts of the matrix: W, R
rnn_layer.multiplier, # 3: four output parts of the matrix for all gates in order: i, o, f, c
rnn_layer.hidden_size, # 4: output size per direction and gate
])
W, R = [x.reshape([
num_directions, # 0: num of directions
rnn_layer.num_layers, # 1: num_layers
rnn_layer.multiplier, # 2: four output parts of the matrix for all gates in order: i, o, f, c
rnn_layer.hidden_size, # 3: output size per direction and gate
-1]) # 4: input size/hidden size in W/R correspondingly
for x in (W, R)]
input_size = match['input'].shape[2]
assert compatible_dims(input_size, W.shape[-1])
# Reorder gates: iofc --> fico
gate_reorder = rnn_layer.gate_order
W, R = (np.take(x, gate_reorder, axis=2) for x in (W, R))
B = np.take(B, gate_reorder, axis=3)
for blob, port in [(W, 1), (R, 2), (B, 3)]:
Op.create_and_connect_input_data_node(
graph,
rnn_layer,
{'value': blob, 'shape': int64_array(blob.shape)},
{'in': port, 'permutation': None}
)
def replace_pattern(graph: Graph, match: dict):
node = match['op']
axis = node.in_port(1).data.get_value()
size_splits = node.in_port(2).data.get_value()
output_shape = sum([node.out_node(port).shape[axis] for port in node.out_nodes()])
if output_shape == node.in_port(0).data.get_shape()[axis]:
return
if not node.has_valid('out_ports_count'):
node['out_ports_count'] = len(size_splits)
Op.normalize_outputs(node)
def copy_input_blobs(op: Node, copy_op: Node):
"""
Function copy input blob data nodes from restored graph to copied one
:param op: Node from restored graph
:param copy_op: Node from copied graph
:return:
"""
for u, d in op.get_sorted_inputs():
if 'bin' in d:
Op.create_and_connect_input_data_node(
copy_op.graph, copy_op, {
'value': op.in_node(d['in']).value,
'shape': op.in_node(d['in']).shape
}, d)
def extract(cls, node):
shapes = node.pb.attr['shapes'].list.shape
tf_types = node.pb.attr['component_types'].list.type
extracted_types = []
for t in tf_types:
extracted_types.append(tf_dtype_extractor(t))
result_shapes = []
for shape_pb in shapes:
shape = shape_pb.dim
if len(shape) == 3:
result_shapes.append(int64_array([1, shape[0].size, shape[1].size, shape[2].size]))
else:
result_shapes.append(int64_array([dim.size for dim in shape]))
Op.update_node_stat(node, {'shapes': result_shapes, 'types': extracted_types})
return cls.enabled
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 = shape_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 extract(cls, node):
shapes = node.pb.attr['output_shapes'].list.shape
tf_types = node.pb.attr['output_types'].list.type
extracted_types = []
for t in tf_types:
extracted_types.append(tf_dtype_extractor(t))
result_shapes = []
for shape_pb in shapes:
result_shapes.append(tf_tensor_shape(shape_pb))
Op.update_node_stat(
node, {
'shapes': result_shapes,
'types': extracted_types,
'out_ports_count': 1
})
return cls.enabled
def extract(cls, node):
narrow_range = node.pb.attr['narrow_range'].b
num_bits = node.pb.attr['num_bits'].i
levels = 2**num_bits - int(narrow_range)
# we prepare this operation to be converted to FakeQuantize op,
# but input reconnection is needed, so we don't set infer function and type attribute
Op.update_node_stat(
node, {
'op': 'FakeQuantWithMinMaxVars',
'levels': levels,
'narrow_range': narrow_range,
'num_bits': num_bits
})
return cls.enabled
def generate_sub_graph(self, graph: Graph, match: SubgraphMatch):
replacement_desc = match.custom_replacement_desc
op = Op.get_op_class_by_name(replacement_desc.op)(
graph, match.custom_replacement_desc.custom_attributes)
op.default_backend_attrs = list(
match.custom_replacement_desc.custom_attributes.keys())
if 'infer' not in op.attrs:
# update IE attrs
op.substitute_ie_attrs(op.attrs)
node = merge_nodes(graph, match.matched_nodes_names(),
replacement_desc.get_inputs_description(),
replacement_desc.get_outputs_description())
node.name = graph.unique_id(op.attrs['type'])
node_attrs = graph.node[node.id]
# copy attributes which are defined in the custom operation
for key in op.attrs.keys():
if key not in ['name', 'op']:
node_attrs[key] = op.attrs[key]
# functions below should return nothing because 'merge_nodes' already created input/output edges
self.input_edges_match = lambda gr, ma, new_sub_graph: dict() # pylint: disable=method-hidden
self.output_edges_match = lambda gr, ma, new_sub_graph: dict() # pylint: disable=method-hidden
else:
node = op.add_node(name=op.attrs['type'] + '_')
node.type = op.attrs['type']
return {'new_node': node}
def split_helper(node: Node, index: int, direction: str, axis: int = 0):
return Op._create_data_node(
node.graph,
name=node.name + '/SplittedBiLSTM/{}/'.format(direction),
attrs={'value': np.take(node.value, [index], axis),
'shape': shape_array(np.take(node.value, [index], axis).shape)}
)
def insert_abs_max(self, model_graph, node, type_stat, node_name,
**kwargs):
axis_const = self.find_axis(node, kwargs.get('granularity'))
if isinstance(axis_const, str):
return (True, node.name)
abs_node = Abs(node.graph, {
"name": f'abs_{node_name}'
}).create_node_with_data([node.out_node(0)]).in_node(0)
max_op = create_op_node_with_second_input(
node.graph, ReduceMax, int64_array(axis_const),
dict(name=f'{type_stat}_{node_name}'))
if node.graph != model_graph:
Op.create_data_node(max_op.graph, max_op, {'shape': [1]})
max_op['fullname'] = reset_node_fullname(node.fullname, max_op.name)
abs_node.out_port(0).connect(max_op.in_port(0))
return self.insert_result(model_graph, node, max_op, type_stat)
def check_init_states(graph: Graph, match: dict):
"""
Check if cell have initial states and create zeros states if not.
And renumber ports for this states.
"""
rnn_cell = match['rnn_layer']
num_directions = 2 if rnn_cell.direction == 'bidirectional' else 1
batch_size = rnn_cell.in_node(0).shape[rnn_cell.batch_dim]
h_init_port = 5
c_init_port = 6
if 2 not in rnn_cell.in_nodes():
h_shape = [num_directions, batch_size,
rnn_cell.hidden_size] # from ONNX spec
h_init = np.full(h_shape, 0, dtype=np.float32)
Op.create_and_connect_input_data_node(
graph, rnn_cell, {
'value': h_init,
'shape': int64_array(h_init.shape)
}, {
'in': h_init_port,
'permutation': None
})
else:
hidden_state_edge = graph.get_edge_data(
rnn_cell.in_node(2).id, rnn_cell.id)
hidden_state_edge[0]['in'] = h_init_port
if rnn_cell.op == 'LSTM':
if 3 not in rnn_cell.in_nodes():
c_shape = [num_directions, batch_size,
rnn_cell.hidden_size] # from ONNX spec
c_init = np.full(c_shape, 0, dtype=np.float32)
Op.create_and_connect_input_data_node(
graph, rnn_cell, {
'value': c_init,
'shape': int64_array(c_init.shape)
}, {
'in': c_init_port,
'permutation': None
})
else:
cell_state_edge = graph.get_edge_data(
rnn_cell.in_node(3).id, rnn_cell.id)
cell_state_edge[0]['in'] = c_init_port
def add_output_reshape(graph: Graph, match: dict):
"""
Since MXNet Y output shape is [batch_size, seq_len, hidden_size * num_directions] we need to add reshape
from above common format [batch_size, num_directions, seq_len, hidden_size] to MXNet format.
"""
lstm = match['rnn_layer']
input = match['input']
if not lstm.has_num_directions:
return
old_data_node = lstm.out_node(0)
num_directions = 2 if lstm.direction in ['bidirectional'] else 1
mxnet_shape = lstm.out_node(0).shape.copy()
if lstm.batch_dim == 0:
mo_shape = shape_array([
input.shape[lstm.batch_dim], input.shape[lstm.sequence_dim],
lstm.hidden_size
])
else:
mo_shape = shape_array([
input.shape[lstm.sequence_dim], input.shape[lstm.batch_dim],
lstm.hidden_size
])
if lstm.has_num_directions:
mo_shape = shape_insert(mo_shape, 1, np.int64(num_directions))
lstm_name = lstm.soft_get('name', lstm.id)
new_data = Op._create_data_node(graph,
name=lstm_name +
'/Data/Reshape_mxnet/',
attrs={'shape': mo_shape})
graph.remove_edge(lstm.id, old_data_node.id)
graph.add_edge(lstm.id, new_data.id, key=0, out=0)
# Add Transpose
permute_order = Const(
graph, {
'name': lstm_name + '/Transpose_mxnet_order',
'value': int64_array([0, 2, 1, 3])
}).create_node_with_data()
permute_data = Transpose(graph, {
'name': lstm_name + '/Transpose_mxnet/'
}).create_node_with_data([new_data, permute_order])
# Add Reshape
reshape = Reshape(graph, {'name': lstm_name + '/Reshape_mxnet/'})
reshape_dim_data = Const(
graph, {
'name': lstm_name + '/Reshape_mxnet_dim',
'value': int64_array(unmask_shape(mxnet_shape))
}).create_node_with_data()
reshape.create_node_with_data([permute_data, reshape_dim_data],
dict(),
data_nodes=[old_data_node])
def infer(node: Node):
# there are limitations coming from ONNX LSTM definition and normalization rules
assert len(node.in_nodes()) >= 3 # X, W and R
assert len(node.in_nodes()) <= 7
assert len(node.out_nodes()) <= 3
assert node.batch_dim <= 1
assert node.sequence_dim <= 1
assert node.batch_dim != node.sequence_dim
assert node.direction in ['forward', 'reverse', 'bidirectional']
if node.blobs_wrb:
mark_input_bins(node, ['W', 'R', 'B'])
else:
mark_input_bins(node)
input_shape = node.in_node(0).shape
assert len(input_shape) == 3
for port in [2, 3]:
if port in node.in_nodes() and len(node.in_node(port).in_nodes()) > 0 and \
'zero_shapes' in node.in_node(port).in_node():
for i in node.in_node(port).in_node().zero_shapes:
if node.in_node(port).shape[i] != input_shape[i]:
node.in_node(port).value = np.repeat(
node.in_node(port).value, input_shape[i], axis=i)
node.in_node(port).shape[i] = input_shape[i]
out_shape = shape_array([
input_shape[node.sequence_dim], input_shape[node.batch_dim],
node.hidden_size
])
assert not node.has_num_directions or node.sequence_dim == 0, \
'If has_num_directions == True, then node.sequence_dim should be equal 0, but it is {}'.format(
node.sequence_dim)
num_directions = 2 if node.direction in ['bidirectional'] else 1
num_layers = node.num_layers
if node.has_num_directions:
# insert extra dimension to output shape for num_directions
out_shape = shape_insert(out_shape, 1, np.int64(num_directions))
node.out_node(0).shape = out_shape
# extra outputs for hidden/cell states
state_size = shape_array([input_shape[1], node.hidden_size])
if node.has_num_directions:
state_size = shape_insert(state_size, 0,
num_directions * num_layers)
for i in [1, 2]:
if i not in node.out_nodes():
data_node = Op._create_data_node(node.graph,
name=node.node +
'/ExtraOutput/' + str(i),
attrs={'executable': True})
node.graph.add_edge(node.id, data_node.id, key=0, out=i)
add_opoutput(node.graph, data_node.id, 0, False)
else:
data_node = node.out_node(i)
data_node.shape = state_size.copy()
def split_helper(node, index: int, direction: str):
return Op._create_data_node(node.graph,
name=node.name +
'/SplittedBiLSTM/{}/'.format(direction),
attrs={
'value':
node.value[index],
'shape':
int64_array(node.value[index].shape)
})
def get_new_cell(multilayer_cell: Node, number: int):
cell_class = Op.get_op_class_by_name(multilayer_cell.op)
new_cell = lambda graph, attrs: cell_class(graph, attrs)
attrs = multilayer_cell.attrs().copy()
new_attrs = {
'num_layers': 1,
'multilayers': False,
'name': multilayer_cell.name + '/LayerSplittedLSTM/{}'.format(number),
}
attrs.update(new_attrs)
return new_cell(multilayer_cell.graph, attrs)
def split_data(self, data: Node):
""" Helper. Split data node into two part along 0 axis """
assert len(data.shape) == 3
assert data.shape[0] == 2
output_data = [Op._create_data_node(data.graph,
name=data.name + '/SplittedBiLSTM/{}'.format(['forward', 'reverse'][i])) for
i in [0, 1]]
split_op = Split(data.graph, dict(name=data.name + '/DecomposedBiLSTM_0', num_splits=2))
axis_const = Const(data.graph, {'name': data.name + '/DecomposedBiLSTM_0' + '/Split_axis',
'value': np.int64(0)}).create_node_with_data()
return split_op.create_node_with_data([data, axis_const], data_nodes=output_data)
def get_new_cell(bidirectional_cell: Node, direction: str):
assert direction in ['forward', 'reverse']
cell_class = Op.get_op_class_by_name(bidirectional_cell.op)
new_cell = lambda graph, attrs: cell_class(graph, attrs)
attrs = bidirectional_cell.attrs().copy()
new_attrs = {
'direction': direction,
'name': bidirectional_cell.name + '/Split/' + direction,
}
attrs.update(new_attrs)
# split bidirectional activations
assert 'activations' in attrs
if attrs['activations'] is not None and len(attrs['activations']) > 1:
assert len(attrs['activations']) == 2, 'Bidirectional RNN should have 2 activations'
activations = attrs['activations']
attrs['activations'] = [activations[0 if direction == 'forward' else 1]]
return new_cell(bidirectional_cell.graph, attrs)
def find_and_replace_pattern(self, graph: Graph):
"""
This function finds all const data nodes that have more that one consumer and then duplicate them
"""
data_nodes = [Node(graph, id) for id in graph.nodes() if Node(graph, id).soft_get('kind') == 'data']
for node in data_nodes:
# Check that node has const values and more than one consumer
if len(node.in_nodes()) and node.in_node().soft_get('type') == 'Const' and len(node.out_nodes()) > 1 and \
node.value is not None:
# Here we delete all edges between base node and it's consumers (except first), and then duplicate this
# node to connect with other consumers
for v, d in node.get_outputs():
out_node = Node(graph, v)
e_attrs = d
graph.remove_edge(node.id, out_node.id)
data = Op.create_input_data_node(graph, "Copy_{}".format(node.id), mo_array(node.value),
graph.node[node.id])
graph.add_edges_from([(data.id, out_node.id, e_attrs)])
def split_bidirectional(self,
bidirectional_cell: Node,
new_init_hiddens: list,
new_init_cells: list,
splitted_W: tuple,
splitted_R: tuple,
splitted_B: tuple):
"""
Split one bidirectional RNNSequence node into 2 one-directional RNNSequence nodes.
All input data nodes should be already prepared; they are
have 2 in the num_dir dimension.
"""
all_outputs = []
for i in [0, 1]:
direction = ['forward', 'reverse'][i]
op = self.get_new_cell(bidirectional_cell, direction)
output_data = Op._create_data_node(
bidirectional_cell.graph,
name=bidirectional_cell.out_node(0).name + '/Split/' + str(i),
attrs={'shape': bidirectional_cell.out_node(0).shape.copy()}
)
assert output_data.shape[1] == 2
output_data.shape[1] = 1
output_hidden = Op._create_data_node(
bidirectional_cell.graph,
name=bidirectional_cell.out_node(1).name + '/Split/' + str(i),
attrs={'shape': bidirectional_cell.out_node(1).shape.copy()}
)
assert output_hidden.shape[0] == 2
output_hidden.shape[0] = 1
data_nodes = [
output_data,
output_hidden,
]
if bidirectional_cell.op == 'LSTM':
output_cell = Op._create_data_node(
bidirectional_cell.graph,
name=bidirectional_cell.out_node(2).name + '/Split/' + str(i),
attrs={'shape': bidirectional_cell.out_node(2).shape.copy()}
)
assert output_cell.shape[0] == 2
output_cell.shape[0] = 1
data_nodes.append(output_cell)
all_outputs.append(
op.create_node_with_data(
inputs=[
bidirectional_cell.in_node(0),
splitted_W[i],
splitted_R[i],
splitted_B[i],
None,
new_init_hiddens[i],
new_init_cells[i] if bidirectional_cell.op == 'LSTM' else None,
],
data_nodes=data_nodes
)
)
return all_outputs
def copy_graph_with_ops(graph: Graph) -> Graph:
"""
Function to copy graph and apply extenders to appropriate nodes
:param graph: Graph to copy
:return:Copied graph with applied extenders
"""
new_graph = Graph()
new_graph.stage = 'back'
new_graph.graph = graph.graph
node_connections = dict()
mapping_of_old_idx_into_new = dict()
restore_correct_ports(graph)
# Nodes preprocessing stage in source graph
# Firstly propagate values only for Const nodes, because other preprocessings
# assumes Const nodes are already preprocessed.
for op in graph.get_op_nodes(type='Const'):
preprocessing_op_nodes[op.type](op)
for op in graph.get_op_nodes():
if op.soft_get('type') != 'Const' and op.soft_get(
'type') in preprocessing_op_nodes:
preprocessing_op_nodes[op.type](op)
# Create a new copy of graph with correct attributes (shape & type infer, backend attrs etc.)
for op in graph.get_op_nodes():
# Save input shapes restored from IR
op['old_input_shapes'] = list()
for n in op.in_nodes():
op.old_input_shapes.append(int64_array(op.in_node(n).shape))
# Apply extenders to nodes in source graph
if op.type in Extender.registered_ops:
Extender.get_extender_class_by_name(op.type).extend(op)
else:
log.debug(
'Extender for node {} with type={} not found, please note.'.
format(op.name, op.type))
# Add node with necessary type and extended attrs in new graph
op_type = op.soft_get('type_to_create', op.type)
if op_type in custom_ops:
node = custom_ops[op_type](new_graph, op.attrs()).create_node()
else:
if op_type not in Op.registered_ops:
log.warning(
'Operation {} is not found in MO operations, please check it! '
'Simple shape infer function is used'.format(op_type))
node = Op(new_graph, op.attrs()).create_node()
assert 'type' in node, 'Operation {} have no `type` attribute.'.format(
node.soft_get('name'))
node['op'] = node.type
node['infer'] = Extender.use_shapes_from_ir
if 'ir_data_attrs' in op:
node['IE'] = [('layer', [
('id', lambda node: node.node), 'name', 'type',
'version'
], [('data', list(op.ir_data_attrs.keys()), []), '@ports',
'@consts'])]
else:
node = Op.get_op_class_by_name(op_type)(
new_graph, op.attrs()).create_node()
# Fill out_ports_count attribute
if 'out_ports_count' not in node and node.soft_get(
'type') != 'Result':
node['out_ports_count'] = len(op.out_edges())
# This attribute is no longer needed and we can delete it
if 'ir_data_attrs' in node:
del node['ir_data_attrs']
if op.has_and_set('need_copy_input_blobs'):
copy_input_blobs(op, node)
# Collect node connections
mapping_of_old_idx_into_new[op.id] = node.id
node_connections[op.id] = collect_node_outputs(op)
# Restore connections in new graph
for input_node_idx, its_outputs in list(node_connections.items()):
for out_port_idx, out_port_dest in its_outputs.items():
for dest_in_port_idx, dest_node_idx in out_port_dest:
src = Node(new_graph,
mapping_of_old_idx_into_new[input_node_idx])
dst = Node(new_graph,
mapping_of_old_idx_into_new[dest_node_idx])
src.out_port(out_port_idx).connect(
dst.in_port(dest_in_port_idx))
# Nodes postprocessing stage in new graph
for op in new_graph.get_op_nodes():
# Call normalize node outputs for restored operations to connect temporary Result operations for disconnected
# output ports. We need to do that for correct shape inference. These Result operations will be removed during
# IR emitting. For TopK operation outputs normalizing we should use specific
# function TopKNormalizer.normalize_outputs.
if op.soft_get('type') != 'TopK':
Op.normalize_outputs(op)
# Set correct_data_type attribute to Const data nodes to correct processing of restored values
if op.soft_get('type') == 'Const':
assert len(op.out_nodes()) == 1 and op.out_node(0).soft_get('kind') == 'data',\
'Const node {} not properly corrected to appropriate data node'.format(op.soft_get('name'))
op.out_node(0)['correct_data_type'] = True
if op.has_and_set('rt_info'):
op.out_node(0)['rt_info'] = op.rt_info
# operations postprocessing with some special types
if op.soft_get('type') in postprocessing_op_nodes:
postprocessing_op_nodes[op.type](op)
restore_tensor_names(op)
# clean up graph to shape inference
new_graph.clean_up()
return new_graph
def extract(cls, node):
Op.update_node_stat(node, {'op': 'FakeConst'})
return cls.enabled
def find_and_replace_pattern(self, graph: Graph):
for split_node in graph.get_op_nodes(op='Split'):
Op.normalize_outputs(split_node)
def extract(cls, node):
attrs = get_attrs(node)
Op.update_node_stat(node, attrs)
return cls.enabled
def find_and_replace_pattern(self, graph: Graph):
# Iterate over all data nodes and find all with >= 1 consumers
for input_data in list(graph.get_data_nodes()):
# We don't use constant data nodes
if input_data.value is not None:
continue
if input_data.shape is None:
continue
input_shape = shape_array(input_data.shape)
# Get all unique StridedSlice consumers
out_nodes = [node for node in input_data.out_nodes() if node.op == 'StridedSlice' and
node.in_node(0).id == input_data.id]
if len(out_nodes) <= 1:
continue
valid_for_replacement = True
for n in out_nodes:
if any(not isinstance(s, slice) for s in n.slices):
# this is a slice with dynamic dimension. Such operation is not valid for replacement
valid_for_replacement = False
if not valid_for_replacement:
continue
sorted_out_nodes = sorted(out_nodes, key=lambda n: list(n.slices))
out_nodes = unique_by(sorted_out_nodes, strided_slices_equality)
for node in out_nodes:
if len(node.slices) != len(out_nodes[0].slices):
valid_for_replacement = False
# Detect dimension for splitting
split_channel_dim = None
for dim_id, s in enumerate(out_nodes[0].slices):
l, r, stride = s.start, s.stop, s.step
# if both l and r are None then the dimension is not sliced
if (l != 0 or r != input_shape[dim_id]) and (l is not None or r is not None):
if split_channel_dim is None:
split_channel_dim = dim_id
else:
valid_for_replacement = False
if split_channel_dim is None:
valid_for_replacement = False
# split_dims contains tuples with split range and output data node
split_dims = []
for out_id, node in enumerate(out_nodes):
# Check that StridedSlice op has stride eq 1 and splits only feature channel
for id, s in enumerate(node.slices):
l, r, stride = s.start, s.stop, s.step
# We don't support StridedSlice with stride != 1
if stride != 1:
valid_for_replacement = False
if id == split_channel_dim:
split_dims.append((s.start, s.stop, node.out_node()))
if not valid_for_replacement:
continue
# Check feature split intersection
final_data_nodes_list = []
sorted_split_dims = sorted(split_dims, key=lambda item: (item[0], item[1]))
# check if we have similar StridedSlice operations with different outputs
prev_sd = sorted_split_dims[0]
to_remove = []
for i in range(1, len(sorted_split_dims)):
if sorted_split_dims[i][0] == prev_sd[0] and sorted_split_dims[i][1] == prev_sd[1] and sorted_split_dims[i][2].name != prev_sd[2].name:
cur_node = sorted_split_dims[i][2]
for out in cur_node.out_nodes():
attrs = deepcopy(graph.get_edge_data(cur_node.id, out.id)[0])
graph.remove_edge(cur_node.id, out.id)
graph.add_edge(prev_sd[2].id, out.id, **attrs)
to_remove.append(i)
for ind in reversed(to_remove):
sorted_split_dims.pop(ind)
size_splits = []
prev_r = 0
for l, r, out in sorted_split_dims:
# Split dims shouldn't intersect
if l < prev_r:
valid_for_replacement = False
prev_r = r
if prev_r > input_shape[split_channel_dim]:
valid_for_replacement = False
if not valid_for_replacement:
continue
prev_r = 0
for l, r, out in sorted_split_dims:
# Save missing tensor part
if l > prev_r:
shape = mo_array(input_shape)
size_splits.append(l - prev_r)
shape[split_channel_dim] = l - prev_r
data_node = Op._create_data_node(graph, 'fake_data_'+out_nodes[0].name, {'shape': shape})
add_opoutput(graph, data_node.id, 0, False, keep_output_port=True)
final_data_nodes_list.append(data_node)
prev_r = r
size_splits.append(r - l)
final_data_nodes_list.append(out)
if prev_r < input_shape[split_channel_dim]:
# Add last part of tensor
shape = input_shape.copy()
shape[split_channel_dim] = input_shape[split_channel_dim] - prev_r
size_splits.append(input_shape[split_channel_dim] - prev_r)
data_node = Op._create_data_node(graph, 'fake_data_'+out_nodes[0].name, {'shape': shape})
add_opoutput(graph, data_node.id, 0, False, keep_output_port=True)
final_data_nodes_list.append(data_node)
for node in out_nodes:
if not np.all([x == 0 for x in node.shrink_axis_mask]):
out_node = node.out_node()
if np.any(node['shrink_axis_mask']):
self.add_squeeze_for_shrink(graph, node)
if np.any(node['new_axis_mask']):
self.add_unsqueeze_for_new(graph, node)
for i in range(len(final_data_nodes_list)):
if final_data_nodes_list[i].name == out_node.name:
final_data_nodes_list[i] = node.out_node()
break
# Insert Split layer and remove old StridedSlice layers
# 1. Remove connections from input_data to StridedSlice ops
out_data_nodes = []
name_for_future_split = out_nodes[0].name
for node in out_nodes:
out_data_nodes.append(node.out_node())
graph.remove_edge(input_data.id, node.id)
graph.remove_edge(node.id, node.out_node().id)
graph.remove_node(node.id)
log.debug("Removed: {}".format(node.id))
# 2. Create Split layer and reorder outputs
name = name_for_future_split + "/Split"
axis_const = Const(graph, {'value': int64_array(split_channel_dim),
'name': name + '/Axis'}).create_node_with_data()
size_splits_const = Const(graph, {'value': int64_array(size_splits),
'name': name + '/Sizes'}).create_node_with_data()
split = VariadicSplit(graph, dict(name=name, out_ports_count=len(size_splits)))
split.create_node_with_data(inputs=[input_data, axis_const, size_splits_const],
data_nodes=final_data_nodes_list)
def repack_weights(graph: Graph, match: dict):
input = match['input']
rnn_layer = match['rnn_layer']
params = match['params'].value.copy()
graph.remove_edge(match['params'].id, rnn_layer.id)
input_size = input.shape[2]
direction = 2 if rnn_layer.has_num_directions else 1
bsize = (2 * rnn_layer.hidden_size * direction *
1) * rnn_layer.multiplier
W = mo_array(params[0:len(params) - bsize])
B = mo_array(params[len(params) - bsize:])
W = W.reshape((direction, -1))
B = B.reshape((direction, -1))
W, R = mo_array(W[:, 0:rnn_layer.hidden_size * rnn_layer.multiplier *
input_size]), mo_array(
W[:, rnn_layer.hidden_size *
rnn_layer.multiplier * input_size:])
W, R = [
x.reshape([
direction, # 0: num of directions
1, # 1: num_cells
rnn_layer.
multiplier, # 2: four output parts of the matrix for all gates
rnn_layer.hidden_size, # 3: output size per direction and gate
-1
]) # 4: input size/hidden size in W/R correspondingly
for x in (W, R)
]
assert W.shape[-1] == input_size
assert R.shape[-1] == rnn_layer.hidden_size
B = B.reshape([
direction, # 0: num of directions, limitation: should be 1
1,
2, # 3: num of component B
rnn_layer.
multiplier, # 1: four output parts of the matrix for all gates in order: i, f, c, o
rnn_layer.hidden_size, # 2: output size per direction and gate
])
# Reorder gates: ifco --> fico
gate_reorder = rnn_layer.gate_order
W = np.take(W, gate_reorder, axis=2)
R = np.take(R, gate_reorder, axis=2)
B = np.take(B, gate_reorder, axis=3)
# Add ports to rnn_layer
rnn_layer.add_sequence_of_ports(type='in', rng=range(7))
for blob, port in [(W, 1), (R, 2), (B, 3)]:
Op.create_and_connect_input_data_node(
graph, rnn_layer, {
'value': blob,
'shape': int64_array(blob.shape)
}, {
'in': port,
'permutation': None
})