218 lines
9.2 KiB
Python
218 lines
9.2 KiB
Python
# Copyright (C) 2018-2021 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import numpy as np
|
|
|
|
from mo.front.common.partial_infer.utils import int64_array
|
|
from mo.front.tf.graph_utils import create_op_node_with_second_input
|
|
from mo.graph.graph import Graph
|
|
from mo.middle.replacement import MiddleReplacementPattern
|
|
from mo.ops.concat import Concat
|
|
from mo.ops.const import Const
|
|
from mo.ops.op import Op
|
|
from mo.ops.reshape import Reshape
|
|
from mo.ops.shape import Shape
|
|
from mo.ops.unsqueeze import Unsqueeze
|
|
from mo.utils.shape import node_to_get_shape_value_of_indices
|
|
|
|
|
|
class RNNSequenceNormalize(MiddleReplacementPattern):
|
|
"""
|
|
This class normalize RNNSequence layers to IE-compatible from of weights, inputs and outputs.
|
|
|
|
In this pass next will be done:
|
|
1. Weights repack (squeeze all useless shapes in all blobs and concatenate W and R together, also add
|
|
bin param and all similar staff )
|
|
1. UNSqueeze num directions (in states and )
|
|
2. Initial states squeeze
|
|
4. Renumbering inputs
|
|
5. Ports checks
|
|
|
|
After this normalization this layer will have next format of inputs:
|
|
0: X input data, shape [batch_size, seq_len, input_size]
|
|
1: WR weights blob, shape [M * hidden_size, hidden_size + input_size]
|
|
2: B biases blob, shape [M * hidden_size]
|
|
3: (optional) sequence_length, shape [batch_size]
|
|
4: initial hidden state, shape [batch_size, hidden_size]
|
|
5: (only for LSTM) initial cell state, shape [batch_size, hidden_size]
|
|
6: (optional for LSTM) Peepholes weights, shape [(M - 1) * hidden_size]
|
|
|
|
"""
|
|
force_shape_inference = True
|
|
|
|
def run_after(self):
|
|
from extensions.middle.DecomposeBidirectionalRNNSequence import DecomposeBidirectionalRNNSequence
|
|
return [DecomposeBidirectionalRNNSequence]
|
|
|
|
def pattern(self):
|
|
return dict(
|
|
nodes=[
|
|
('rnn_layer', dict(kind='op', type='RNNSequence')),
|
|
('input', dict(kind='data')),
|
|
('W', dict(kind='data')),
|
|
('R', dict(kind='data')),
|
|
('B', dict(kind='data')),
|
|
],
|
|
edges=[
|
|
('input', 'rnn_layer', {'in': 0}),
|
|
('W', 'rnn_layer', {'in': 1}),
|
|
('R', 'rnn_layer', {'in': 2}),
|
|
('B', 'rnn_layer', {'in': 3}),
|
|
],
|
|
)
|
|
|
|
def replace_pattern(self, graph: Graph, match: dict):
|
|
self.repack_weights(graph, match)
|
|
if match['rnn_layer'].has_num_directions:
|
|
self.unsqueeze_num_directions(graph, match)
|
|
self.squeeze_initial_states(graph, match)
|
|
self.reordering_inputs(graph, match)
|
|
# some additional checks for ports number and similar stuff
|
|
|
|
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}
|
|
)
|
|
|
|
@staticmethod
|
|
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])
|
|
@staticmethod
|
|
def squeeze_initial_states(graph: Graph, match: dict):
|
|
"""
|
|
Squeeze input initial states of recurrent node to 2-D shape.
|
|
"""
|
|
hidden_init_port = 5
|
|
cell_init_port = 6
|
|
|
|
rnn_layer = match['rnn_layer']
|
|
# Add input ports to rnn_layer
|
|
rnn_layer.add_sequence_of_ports(type='in', rng=range(7))
|
|
rnn_layer_name = rnn_layer.soft_get('name', rnn_layer.id)
|
|
|
|
assert hidden_init_port in rnn_layer.in_nodes()
|
|
hidden_size = rnn_layer.hidden_size
|
|
shape = Shape(graph, dict(name=rnn_layer_name + '/ShapeOf')).create_node()
|
|
rnn_layer.in_port(0).get_source().connect(shape.in_port(0))
|
|
|
|
batch = node_to_get_shape_value_of_indices(shape, int64_array([rnn_layer.batch_dim]))
|
|
new_dim = create_op_node_with_second_input(graph, Concat, second_input_value=int64_array([hidden_size]),
|
|
op_attrs=dict(name=rnn_layer_name + '/HiddenStateResizeDim',
|
|
in_ports_count=2, axis=0), input_node=batch)
|
|
reshape_h = Reshape(graph, dict(name=rnn_layer_name + '/HiddenStateResize', override_output_shape=True)).create_node()
|
|
new_dim.out_port(0).connect(reshape_h.in_port(1))
|
|
rnn_layer.in_port(hidden_init_port).get_connection().insert_node(reshape_h)
|
|
|
|
if rnn_layer.op == 'LSTM':
|
|
assert cell_init_port in rnn_layer.in_nodes()
|
|
|
|
reshape_c = Reshape(graph, dict(name=rnn_layer_name + '/CellStateResize', override_output_shape=True)).create_node()
|
|
new_dim.out_port(0).connect(reshape_c.in_port(1))
|
|
rnn_layer.in_port(cell_init_port).get_connection().insert_node(reshape_c)
|
|
|
|
@staticmethod
|
|
def reordering_inputs(graph: Graph, match: dict):
|
|
"""
|
|
Reorder (renumbering) inputs to described format. We need to renumber initial states ports.
|
|
"""
|
|
rnn_layer = match['rnn_layer']
|
|
assert 5 in rnn_layer.in_nodes()
|
|
hidden_state_edge = graph.get_edge_data(rnn_layer.in_node(5).id, rnn_layer.id)
|
|
hidden_state_edge[0]['in'] = 4
|
|
|
|
if rnn_layer.op == 'LSTM':
|
|
assert 6 in rnn_layer.in_nodes()
|
|
cell_state_edge = graph.get_edge_data(rnn_layer.in_node(6).id, rnn_layer.id)
|
|
cell_state_edge[0]['in'] = 5
|
|
|
|
@staticmethod
|
|
def ports_checks(graph: Graph, match: dict):
|
|
"""
|
|
Check that all mandatory ports is present.
|
|
"""
|
|
rnn_layer = match['rnn_layer']
|
|
mandatory_ports = [0, 1, 2, 4]
|
|
|
|
if rnn_layer.op == 'LSTM':
|
|
mandatory_ports.append(5)
|
|
|
|
assert set(rnn_layer.in_nodes().keys()) >= set(mandatory_ports) |