[CFE] Import TensorFlow RNN/LSTM/GRU

seanshpark commented 1 year ago

Let's findout how to import TensorFlow LSTM operator.

reference:

seanshpark commented 1 year ago

related #8217 previous #4528

derived tasks from study

[x] study LSTM
[x] unroll LSTM
[x] codegen unrolled LSTM for TRIX
[x] generat tflite fused UnidirectionalSequenceLSTM Op
[x] unroll GRU
[x] codegen unrolled GRU for TRIX --> fails
[x] unroll SimpleRNN
[x] codegen unrolled SimpleRNN for TRIX
[x] RNN + SimpleRNNCell
[x] RNN + LSTMCell
[x] RNN + GRUCell

seanshpark commented 1 year ago

With https://github.com/Samsung/ONE/blob/6a8a5e3a7247a246271286ada127a2eba6f8c4ef/res/TensorFlowPythonExamples/examples/lstm/__init__.py + #9899

        # tflite export for experiments
        converter = tf.lite.TFLiteConverter.from_keras_model(m.model)
        converter.allow_custom_ops = True
        converter.experimental_new_converter = True
        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]

        tflite_model = converter.convert()

gives

error: failed to legalize operation 'tf.TensorListReserve' that was explicitly marked illegal error: Lowering tensor list ops is failed. Please consider using Select TF ops and disabling _experimental_lower_tensor_list_ops flag in the TFLite converter object. For example, converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]\n converter._experimental_lower_tensor_list_ops = False

seanshpark commented 1 year ago

With converter._experimental_lower_tensor_list_ops = False,

The following operation(s) need TFLite custom op implementation(s): Custom ops: TensorListFromTensor, TensorListGetItem, TensorListReserve, TensorListSetItem, TensorListStack Details: tf.TensorListFromTensor(tensor<4x?x4xf32>, tensor<2xi32>) -> (tensor<!tf_type.variant<tensor<?x4xf32>>>) : {device = ""} tf.TensorListGetItem(tensor<!tf_type.variant<tensor<?x4xf32>>>, tensor, tensor<2xi32>) -> (tensor<?x4xf32>) : {device = ""} tf.TensorListReserve(tensor<2xi32>, tensor) -> (tensor<!tf_type.variant<tensor<?x2xf32>>>) : {device = ""} tf.TensorListSetItem(tensor<!tf_type.variant<tensor<?x2xf32>>>, tensor, tensor<?x2xf32>) -> (tensor<!tf_type.variant<tensor<?x2xf32>>>) : {device = ""} tf.TensorListStack(tensor<!tf_type.variant<tensor<?x2xf32>>>, tensor<2xi32>) -> (tensor<?x?x2xf32>) : {device = "", num_elements = -1 : i64} See instructions: https://www.tensorflow.org/lite/guide/ops_custom

seanshpark commented 1 year ago

Custom ops: TensorListFromTensor, TensorListGetItem, TensorListReserve, TensorListSetItem, TensorListStack

NOT feasible... or is it?

seanshpark commented 1 year ago

keras.layers.LSTM has option unroll, default False. Let's give this a True and see...

seanshpark commented 1 year ago

--> produces an unrolled tflite

seanshpark commented 1 year ago

--> there are unknown shape

Operands: T(subgraph index : tensor index) TYPE (shape) (shape_signature) B(buffer index) (variable) OperandName
T(0:0) FLOAT32 (1, 4, 4) (-1, 4, 4) B(1) serving_default_lstm_input:0
...
T(0:12) FLOAT32 (4, 1, 4) (4, -1, 4) B(13) sequential/lstm/transpose

seanshpark commented 1 year ago

tried...

import tensorflow as tf
import numpy as np
from tensorflow import keras

model = keras.Sequential()
shape = (4, 4)

model.add(keras.layers.LSTM(2, input_shape=shape, unroll=True))
model.compile(optimizer="adam", loss="mse", metrics=["accuracy"])

input_data = np.random.random((1, 4, 4))
target_data = np.random.random((1, 2))
model.fit(input_data, target_data, batch_size=1)

--> no luck

seanshpark commented 1 year ago

give input with fixed shape?

import tensorflow as tf
import numpy as np
from tensorflow import keras

model = keras.Sequential()
shape = (4, 4)

model.add(keras.layers.InputLayer(input_shape=shape, batch_size=1))
model.add(keras.layers.LSTM(2, input_shape=shape, unroll=True))

--> add InputLayer with fixed batch size. --> Yes! --> also TRIX compilation is also OK

[onecc]
one-import-tf=True
one-import-tflite=False
one-import-bcq=False
one-import-onnx=False
one-optimize=True
one-quantize=True
one-pack=False
one-codegen=False

[one-import-tf]
input_path=lstm_unroll.h5
output_path=lstm_unroll.circle
model_format=keras_model

[one-optimize]
input_path=./lstm_unroll.circle
output_path=./lstm_unroll.opt.circle

[one-quantize]
input_path=./lstm_unroll.opt.circle
output_path=./lstm_unroll.opt.q8.circle

seanshpark commented 1 year ago

lstm_unroll.opt.q8.circle.zip

seanshpark commented 1 year ago

Without unroll=True argument,

# NOTE tested with TF 2.8.0
import tensorflow as tf
import numpy as np

from tensorflow import keras

model = keras.Sequential()
shape = (4, 4)

model.add(keras.layers.InputLayer(input_shape=shape, batch_size=1))
model.add(keras.layers.LSTM(2, input_shape=shape))

gives tflite with one UnidirectionalSequenceLSTM Op.

seanshpark commented 1 year ago

And with return_sequences=True

# NOTE tested with TF 2.8.0
import tensorflow as tf
import numpy as np

from tensorflow import keras

model = keras.Sequential()
shape = (4, 4)

model.add(keras.layers.InputLayer(input_shape=shape, batch_size=1))
model.add(keras.layers.LSTM(2, input_shape=shape, return_sequences=True))

gives single Op.

seanshpark commented 1 year ago

from https://www.tensorflow.org/lite/models/convert/rnn#bidirectional_lstm,

says This is future work

tried this like


# NOTE tested with TF 2.8.0
import tensorflow as tf
import numpy as np

from tensorflow import keras

model = keras.Sequential() shape = (4, 4)

model.add(keras.layers.InputLayer(input_shape=shape, batch_size=1)) lstmf = keras.layers.LSTM(2, return_sequences=True) lstmb = keras.layers.LSTM(2, return_sequences=True, go_backwards=True) model.add(keras.layers.Bidirectional(lstmf, backward_layer=lstmb, input_shape=shape))


got
![image](https://user-images.githubusercontent.com/4616940/197650446-ad492d9a-9033-4807-baf5-d32340f7a91f.png)

seanshpark commented 1 year ago

SimpleRNN with unroll=True

# NOTE tested with TF 2.8.0
import tensorflow as tf
import numpy as np

from tensorflow import keras

model = keras.Sequential()
shape = (4, 4)

model.add(keras.layers.InputLayer(input_shape=shape, batch_size=1))
model.add(keras.layers.SimpleRNN(2, input_shape=shape, unroll=True))

gives

TRIX code generaion is OK

seanshpark commented 1 year ago

GRU unroll

# NOTE tested with TF 2.8.0
from tensorflow import keras

model = keras.Sequential()
shape = (4, 4)

model.add(keras.layers.InputLayer(input_shape=shape, batch_size=1))
model.add(keras.layers.GRU(2, input_shape=shape, unroll=True))

seanshpark commented 1 year ago

RNN + LSTMCell

# NOTE tested with TF 2.8.0
from tensorflow import keras

model = keras.Sequential()
shape = (4, 4)

model.add(keras.layers.InputLayer(input_shape=shape, batch_size=1))
lstmcell = keras.layers.LSTMCell(2)
model.add(keras.layers.RNN(lstmcell, input_shape=shape, unroll=True))

TRIX codegen is OK

seanshpark commented 1 year ago

RNN + GRUCell

# NOTE tested with TF 2.8.0
from tensorflow import keras

model = keras.Sequential()
shape = (4, 4)

model.add(keras.layers.InputLayer(input_shape=shape, batch_size=1))
grucell = keras.layers.GRUCell(2)
model.add(keras.layers.RNN(grucell, input_shape=shape, unroll=True))

seanshpark commented 1 year ago

Current CFE UnidirectionalSequenceLSTM status

Op enable issue: #4151
tflchef, luci is OK

Items left to enable

luci-interpreter
value test

seanshpark commented 1 year ago

Enable luci-interpreter

(1) Prepare empty classes

XXXParams in compiler/luci-interpreter/src/core/KernelParams.h
kernel in compiler/luci-interpreter/src/kernels/
loader in compiler/luci-interpreter/src/loader/nodes/

(2) fill in with kernels, tests (3) fix tests to success

seanshpark commented 1 year ago

RNN? GRU? in tflite?

BuiltinOperator_RNN
BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN
BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN
there are no GRU

How can we drop RNN Op from python?

looking inside tensorflow/compiler/mlir/lite...

struct LegalizeUnidirectionalSequenceRnn : public RewritePattern

    if (op->getNumOperands() != 5) {
      op->emitError()
          << "We're expecting 5 inputs for UnidirectionalSequenceRNN, only "
          << op->getNumOperands() << " provided";
      return failure();
    }

    if (op->getNumResults() != 2) {
      op->emitError()
          << "We're expecting 2 inputs for UnidirectionalSequenceRNN, only "
          << op->getNumResults() << " found";
      return failure();
    }

?

in tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc

const char kUnidirectionalSequenceRnnOp[] =
    "name: 'UnidirectionalSequenceRnn' input_arg: {name: 'Input' type: "
    "DT_FLOAT} input_arg: { name: 'Weights' type: DT_FLOAT } "
    "input_arg: { name: 'RecurrentWeights' type: DT_FLOAT } input_arg: { "
    "name: 'Bias' type: DT_FLOAT} "
    "input_arg: { name: 'HiddenState' type: DT_FLOAT} "
    "output_arg: { name: "
    "'LastState' type: DT_FLOAT } output_arg: { name: 'Output' type: "
    "DT_FLOAT} "
    "attr : { name: '_tflite_input_indices' type: 'list(int)'}";

--> how to make this Op(?)/Layer(?) is a question...

SlavikMIPT commented 1 year ago

from https://www.tensorflow.org/lite/models/convert/rnn#bidirectional_lstm,

says This is future work
tried this like

# NOTE tested with TF 2.8.0
import tensorflow as tf
import numpy as np

from tensorflow import keras

model = keras.Sequential()
shape = (4, 4)

model.add(keras.layers.InputLayer(input_shape=shape, batch_size=1))
lstmf = keras.layers.LSTM(2, return_sequences=True)
lstmb = keras.layers.LSTM(2, return_sequences=True, go_backwards=True)
model.add(keras.layers.Bidirectional(lstmf, backward_layer=lstmb, input_shape=shape))

got

According to https://www.tensorflow.org/lite/models/convert/rnn#user-defined_lstm_conversion_examples

Bidirectional LSTM is currently modelled as two UnidirectionalSequenceLSTM operations in TensorFlow Lite. This will be replaced with a single BidirectionalSequenceLSTM op.

Also I studied tensorflow repo and there is recent activity on this. So I am going to implement BidirectionalSequenceLSTM Kernel in ONE - so we can activate this as soon as it will be implemented in tensorflow's converter

Samsung / ONE

[CFE] Import TensorFlow RNN/LSTM/GRU #9895

Enable luci-interpreter