92 lines
3.2 KiB
Python
92 lines
3.2 KiB
Python
# Copyright (C) 2018-2021 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import numpy as np
|
|
|
|
from mo.graph.graph import Graph
|
|
from mo.middle.replacement import MiddleReplacementPattern
|
|
|
|
|
|
class TensorFlowLSTMtoGeneric(MiddleReplacementPattern):
|
|
"""
|
|
Resolves all differences in TensorFlow LSTMCell and Inference Engine LSTMCell:
|
|
- weights transposing
|
|
- shift_const value addition to biases
|
|
- extra inputs deletion
|
|
"""
|
|
enabled = True
|
|
|
|
def run_after(self):
|
|
from extensions.middle.pass_separator import MiddleStart
|
|
return [MiddleStart]
|
|
|
|
def run_before(self):
|
|
from extensions.middle.permute_tensor_iterator import TransposeTensorIteratorLSTM
|
|
return [TransposeTensorIteratorLSTM]
|
|
|
|
|
|
def pattern(self):
|
|
return dict(
|
|
nodes=[('lstm', dict(op='LSTMCell', tf=True))],
|
|
edges=[]
|
|
)
|
|
|
|
def replace_pattern(self, graph: Graph, match: dict):
|
|
weights_node = match['lstm'].in_node(3)
|
|
biases_node = match['lstm'].in_node(4)
|
|
node = match['lstm']
|
|
shift_const = node.shift_const
|
|
|
|
# make sure that the node is the only consumer or weights and biases
|
|
# to let us modify them without hassle
|
|
assert len(weights_node.out_nodes()) == 1
|
|
assert len(biases_node.out_nodes()) == 1
|
|
|
|
# Assign temporary shape for them for easier manipulation
|
|
# TF stores weights in IO order
|
|
input_size = node.in_node(0).shape[1]
|
|
hidden_size = node.in_node(1).shape[1]
|
|
weights = weights_node.value
|
|
biases = biases_node.value
|
|
assert weights.shape[0] == input_size + hidden_size, \
|
|
"weights.shape={} input_size={} hidden_size={}".format(weights.shape, input_size, hidden_size)
|
|
assert weights.shape[1] == biases.shape[0] == 4 * hidden_size, \
|
|
"weights.shape={} biases.shape={} hidden_size={}".format(weights.shape, biases.shape, hidden_size)
|
|
|
|
weights = weights.reshape([
|
|
weights.shape[0],
|
|
4, # gates
|
|
hidden_size
|
|
])
|
|
|
|
biases = biases.reshape([
|
|
4, # gates
|
|
hidden_size
|
|
])
|
|
|
|
# Reorder gates icfo --> fico for both weights and biases
|
|
gate_reorder = [2, 0, 1, 3]
|
|
weights = np.take(weights, gate_reorder, axis=1)
|
|
biases = np.take(biases, gate_reorder, axis=0)
|
|
|
|
# shift_const.value should be added to the first 1/4th part of the biases (f-gate: 0)
|
|
# Note: in case of moving this code up before gate reordering, the addition
|
|
# should be applied at different place
|
|
biases[0] += shift_const
|
|
|
|
# Return to the original shapes
|
|
weights = weights.reshape([weights.shape[0], -1])
|
|
biases = biases.flatten()
|
|
|
|
# TF stores weights in IO, but IE requires it in OI: transpose
|
|
weights = weights.transpose()
|
|
|
|
weights_node.value = weights
|
|
weights_node.shape = np.array(weights.shape, dtype=np.int64)
|
|
biases_node.value = biases
|
|
biases_node.shape = np.array(biases.shape, dtype=np.int64)
|
|
|
|
# Cut all extra inputs off
|
|
for i in range(len(node.inputs), len(node.inputs) + len(node.extra_inputs)):
|
|
node.graph.remove_edge(node.in_node(i).id, node.id)
|