model-optimizer/extensions/middle/decompose_bi_lstm.py

   1 """
   2  Copyright (c) 2018 Intel Corporation
   3
   4  Licensed under the Apache License, Version 2.0 (the "License");
   5  you may not use this file except in compliance with the License.
   6  You may obtain a copy of the License at
   7
   8       http://www.apache.org/licenses/LICENSE-2.0
   9
  10  Unless required by applicable law or agreed to in writing, software
  11  distributed under the License is distributed on an "AS IS" BASIS,
  12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  See the License for the specific language governing permissions and
  14  limitations under the License.
  15 """
  16
  17 import networkx as nx
  18 import numpy as np
  19 from copy import deepcopy
  20
  21 from extensions.ops.lstm_sequence import LSTMSequence
  22 from mo.utils.error import Error
  23 from mo.middle.replacement import MiddleReplacementPattern
  24 from mo.ops.concat import Concat
  25 from mo.ops.op import Op
  26 from mo.ops.split import Split
  27 from mo.graph.graph import Node
  28
  29
  30 class DecomposeBiLSTM(MiddleReplacementPattern):
  31     ''' Decomposes bidirectional LSTMSequence to forward and reverse LSTM ops.
  32
  33         To extract forward and reverse parts from initial blobs, the helper
  34         functions used that should be already built-in into the operation attributes.
  35
  36         Both initial state are split to two part, two parts of the results are concatenated.
  37         Axis of split/concat is completelly defined by ONNX/LSTM specification.
  38     '''
  39
  40     enabled = True
  41
  42     def pattern(self):
  43         return dict(
  44             nodes=[
  45                 ('lstm', dict(kind='op', op='LSTMSequence', format='onnx', direction='bidirectional')),
  46                 ('input', dict(kind='data')),
  47                 ('W', dict(kind='data')),
  48                 ('R', dict(kind='data')),
  49             ],
  50             edges=[
  51                 ('input', 'lstm', {'in': 0}),
  52                 ('W', 'lstm', {'bin': 'W'}),
  53                 ('R', 'lstm', {'bin': 'R'}),
  54             ]
  55         )
  56
  57
  58     def replace_pattern(self, graph: nx.MultiDiGraph, match: dict):
  59         bilstm = match['lstm']
  60         new_init_hiddens = self.split_data(bilstm.in_node(5))
  61         new_init_cells = self.split_data(bilstm.in_node(6))
  62         assert bilstm.has_valid('blob_bidirectional_split'), \
  63             'Node {} doesnt\'t have blob_bidirectional_split attribute defined.'.format(bilstm.soft_get('name'))
  64         splitted_W = bilstm.blob_bidirectional_split(bilstm.in_node(1))
  65         splitted_R = bilstm.blob_bidirectional_split(bilstm.in_node(2))
  66         splitted_B = bilstm.blob_bidirectional_split(bilstm.in_node(3)) if 3 in bilstm.in_nodes() else (None, None)
  67
  68         outputs = self.split_bilstm(
  69             bilstm,
  70             new_init_hiddens,
  71             new_init_cells,
  72             splitted_W,
  73             splitted_R,
  74             splitted_B,
  75         )
  76
  77         self.concat(bilstm, outputs[0], outputs[1], bilstm.out_nodes())
  78
  79     def split_data(self, data: Node):
  80         """ Split data node into two part along 0 axis """
  81         assert len(data.shape) == 3
  82         assert data.shape[0] == 2
  83
  84         output_data = [Op._create_data_node(data.graph, name=data.name + '/SplittedBiLSTM/{}'.format(['forward', 'reverse'][i])) for i in [0, 1]]
  85         split_op = Split(data.graph, dict(name=data.name + '/DecomposedBiLSTM_0', axis=0, num_split=2))
  86         return split_op.create_node_with_data([data], data_nodes=output_data)
  87
  88
  89     def split_bilstm(self,
  90                      bilstm,
  91                      new_init_hiddens,
  92                      new_init_cells,
  93                      splitted_W,
  94                      splitted_R,
  95                      splitted_B):
  96         """ Split one bilstm node into 2 one-directional lstm nodes.
  97
  98             All input data nodes should be already prepared; they are
  99             have 2 in the major dimension.
 100         """
 101         assert len(bilstm.out_nodes()) == 3
 102         all_outputs = []
 103         for i in [0, 1]:
 104             direction = ['forward', 'reverse'][i]
 105             op = LSTMSequence(bilstm.graph, {
 106                 'hidden_size': bilstm.hidden_size,
 107                 'direction': direction,
 108                 'batch_dim': bilstm.batch_dim,
 109                 'sequence_dim': bilstm.sequence_dim,
 110                 'blobs_wrb': bilstm.blobs_wrb,
 111                 'has_num_directions': bilstm.has_num_directions,
 112                 'format': bilstm.format,
 113                 'name': bilstm.name + '/Split/' + direction,
 114             })
 115
 116             output_data = Op._create_data_node(
 117                 bilstm.graph,
 118                 name=bilstm.out_node(0).name + '/Split/' + str(i),
 119                 attrs = {'shape': bilstm.out_node(0).shape.copy()}
 120             )
 121
 122             assert output_data.shape[1] == 2
 123             output_data.shape[1] = 1
 124
 125             output_hidden = Op._create_data_node(
 126                 bilstm.graph,
 127                 name=bilstm.out_node(1).name + '/Split/' + str(i),
 128                 attrs = {'shape': bilstm.out_node(1).shape.copy()}
 129             )
 130
 131             assert output_hidden.shape[0] == 2
 132             output_hidden.shape[0] = 1
 133
 134             output_cell = Op._create_data_node(
 135                 bilstm.graph,
 136                 name=bilstm.out_node(2).name + '/Split/' + str(i),
 137                 attrs = {'shape': bilstm.out_node(2).shape.copy()}
 138             )
 139
 140             assert output_cell.shape[0] == 2
 141             output_cell.shape[0] = 1
 142
 143             all_outputs.append(
 144                 op.create_node_with_data(
 145                     inputs = [
 146                         bilstm.in_node(0),
 147                         splitted_W[i],
 148                         splitted_R[i],
 149                         splitted_B[i],
 150                         None,
 151                         new_init_hiddens[i],
 152                         new_init_cells[i],
 153                     ],
 154                     data_nodes = [
 155                         output_data,
 156                         output_hidden,
 157                         output_cell
 158                     ]
 159                 )
 160             )
 161         return all_outputs
 162
 163
 164     def concat(self, bilstm, forward_outputs, reverse_outputs, final_outputs):
 165         """ Concatenates two set of outputs from BiLSTM """
 166
 167         concat_ops = [
 168             Concat(bilstm.graph, {
 169                 'name': bilstm.name + '/FinalConcat/Data',
 170                 'axis': 1
 171             }),
 172             Concat(bilstm.graph, {
 173                 'name': bilstm.name + '/FinalConcat/HiddenState',
 174                 'axis': 0
 175             }),
 176             Concat(bilstm.graph, {
 177                 'name': bilstm.name + '/FinalConcat/CellState',
 178                 'axis': 0
 179             })
 180         ]
 181
 182         bilstm.graph.remove_node(bilstm.id)
 183
 184         for i in final_outputs:
 185             concat_ops[i].create_node_with_data(
 186                 [forward_outputs[i], reverse_outputs[i]],
 187                 data_nodes=[final_outputs[i]]
 188             )