search

google-ai-edge / LiteRT

LiteRT is the new name for TensorFlow Lite (TFLite). While the name is new, it's still the same trusted, high-performance runtime for on-device AI, now with an expanded vision.
https://ai.google.dev/edge/litert
Apache License 2.0
168 stars 13 forks source link

Native irfft (and inverse_stft) ops in TFLite #130

Open pkgoogle opened 1 day ago

pkgoogle commented 1 day ago

Original Issue: https://github.com/tensorflow/tensorflow/issues/54501 Opening on behalf of @francoisgermain

System information

Describe the feature and the current behavior/state. Now that TFLite supports natively ops for tf.signal.rfft and tf.signal.stft (see #27030), it would be really great to close the loop and add support for their inverse functions tf.signal.irfft and tf.signal.inverse_stft to support more broadly 1-dimensional data models.

Right now, here is what happens with this MWE:

import tensorflow as tf
import tensorflow.keras as tfk
import tensorflow.keras.layers as tfkl

N = 128
fft_length = 1024

input = tfk.Input((N, fft_length // 2 + 1), dtype=tf.dtypes.complex64)
output = tfkl.Lambda(lambda x: tf.signal.irfft(x), output_shape=(N, fft_length))(input)

model = tfk.Model(inputs=input, outputs=output)

tf.get_logger().warning("Begin Tensorflow test...")
test_input = tf.complex(
    tf.random.normal((1, N, fft_length // 2 + 1), dtype=tf.dtypes.float32),
    tf.random.normal((1, N, fft_length // 2 + 1), dtype=tf.dtypes.float32))    
test_output = model(test_input)
tf.get_logger().warning("Tensorflow test finished...")

tf.get_logger().warning("Begin TFLite conversion...")
model.save('./tflite')
converter = tf.lite.TFLiteConverter.from_saved_model('./tflite')
tflite_model = converter.convert()
with open(os.path.join('./tflite', 'model.tflite'), "wb") as f:
    f.write(tflite_model)
tf.get_logger().warning("TFLite conversion finished...")

Output:

2022-02-23 17:05:01.840430: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:936] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-02-23 17:05:02.447990: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:936] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-02-23 17:05:02.448941: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:936] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-02-23 17:05:02.450203: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-02-23 17:05:02.450930: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:936] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-02-23 17:05:02.451762: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:936] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-02-23 17:05:02.452562: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:936] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-02-23 17:05:03.149397: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:936] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-02-23 17:05:03.150348: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:936] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-02-23 17:05:03.151175: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:936] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-02-23 17:05:03.152004: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 10794 MB memory:  -> device: 0, name: Tesla K80, pci bus id: 0000:00:1e.0, compute capability: 3.7
WARNING:tensorflow:Begin Tensorflow test...
WARNING:tensorflow:Tensorflow test finished...
WARNING:tensorflow:Begin TFLite conversion...
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
2022-02-23 17:05:08.281960: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
2022-02-23 17:05:08.542623: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:357] Ignored output_format.
2022-02-23 17:05:08.542684: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:360] Ignored drop_control_dependency.
2022-02-23 17:05:08.544198: I tensorflow/cc/saved_model/reader.cc:43] Reading SavedModel from: ./tflite
2022-02-23 17:05:08.546087: I tensorflow/cc/saved_model/reader.cc:78] Reading meta graph with tags { serve }
2022-02-23 17:05:08.546141: I tensorflow/cc/saved_model/reader.cc:119] Reading SavedModel debug info (if present) from: ./tflite
2022-02-23 17:05:08.549053: I tensorflow/cc/saved_model/loader.cc:228] Restoring SavedModel bundle.
2022-02-23 17:05:08.587791: I tensorflow/cc/saved_model/loader.cc:212] Running initialization op on SavedModel bundle at path: ./tflite
2022-02-23 17:05:08.594120: I tensorflow/cc/saved_model/loader.cc:301] SavedModel load for tags { serve }; Status: success: OK. Took 49927 microseconds.
2022-02-23 17:05:08.602500: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:237] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
loc(callsite(callsite(fused["IRFFT:", "model/lambda/irfft@__inference__wrapped_model_39"] at fused["PartitionedCall:", "PartitionedCall@__inference_signature_wrapper_138"]) at fused["PartitionedCall:", "PartitionedCall"])): error: 'tf.IRFFT' op is neither a custom op nor a flex op
error: failed while converting: 'main':
Some ops are not supported by the native TFLite runtime, you can enable TF kernels fallback using TF Select. See instructions: https://www.tensorflow.org/lite/guide/ops_select
TF Select ops: IRFFT
Details:
    tf.IRFFT(tensor<?x128x513xcomplex<f32>>, tensor<1xi32>) -> (tensor<?x128x1024xf32>) : {device = ""}

Traceback (most recent call last):
  File "/fsx/findsimilar-experiments/tflite_irfft_test.py", line 24, in <module>
    tflite_model = converter.convert()
  File "/fsx/findsimilar-experiments/venv-newtf28/lib/python3.9/site-packages/tensorflow/lite/python/lite.py", line 803, in wrapper
    return self._convert_and_export_metrics(convert_func, *args, **kwargs)
  File "/fsx/findsimilar-experiments/venv-newtf28/lib/python3.9/site-packages/tensorflow/lite/python/lite.py", line 789, in _convert_and_export_metrics
    result = convert_func(self, *args, **kwargs)
  File "/fsx/findsimilar-experiments/venv-newtf28/lib/python3.9/site-packages/tensorflow/lite/python/lite.py", line 1084, in convert
    return self._convert_from_saved_model(graph_def)
  File "/fsx/findsimilar-experiments/venv-newtf28/lib/python3.9/site-packages/tensorflow/lite/python/lite.py", line 967, in _convert_from_saved_model
    result = _convert_saved_model(**converter_kwargs)
  File "/fsx/findsimilar-experiments/venv-newtf28/lib/python3.9/site-packages/tensorflow/lite/python/convert_phase.py", line 213, in wrapper
    raise converter_error from None  # Re-throws the exception.
  File "/fsx/findsimilar-experiments/venv-newtf28/lib/python3.9/site-packages/tensorflow/lite/python/convert_phase.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/fsx/findsimilar-experiments/venv-newtf28/lib/python3.9/site-packages/tensorflow/lite/python/convert.py", line 789, in convert_saved_model
    data = convert(
  File "/fsx/findsimilar-experiments/venv-newtf28/lib/python3.9/site-packages/tensorflow/lite/python/convert.py", line 306, in convert
    raise converter_error
tensorflow.lite.python.convert_phase.ConverterError: <unknown>:0: error: loc(callsite(callsite(fused["IRFFT:", "model/lambda/irfft@__inference__wrapped_model_39"] at fused["PartitionedCall:", "PartitionedCall@__inference_signature_wrapper_138"]) at fused["PartitionedCall:", "PartitionedCall"])): 'tf.IRFFT' op is neither a custom op nor a flex op
<unknown>:0: note: loc(fused["PartitionedCall:", "PartitionedCall"]): called from
<unknown>:0: note: loc(callsite(callsite(fused["IRFFT:", "model/lambda/irfft@__inference__wrapped_model_39"] at fused["PartitionedCall:", "PartitionedCall@__inference_signature_wrapper_138"]) at fused["PartitionedCall:", "PartitionedCall"])): Error code: ERROR_NEEDS_FLEX_OPS
<unknown>:0: error: failed while converting: 'main':
Some ops are not supported by the native TFLite runtime, you can enable TF kernels fallback using TF Select. See instructions: https://www.tensorflow.org/lite/guide/ops_select
TF Select ops: IRFFT
Details:
    tf.IRFFT(tensor<?x128x513xcomplex<f32>>, tensor<1xi32>) -> (tensor<?x128x1024xf32>) : {device = ""}

The TF Select workaround mentioned in the error message does seem to work, so the community may have a workaround in the meantime (with all the caveats mentioned in https://www.tensorflow.org/lite/guide/ops_select), but I think it'd be in the best interest of the community to have the inverse operators supported natively whenever possible, for a much facilitated deployment of models.

Will this change the current api? How?

Not that I can think of, but maybe I don't fully understand the question.

Who will benefit with this feature?

Every TensorFlow user needing to deploy a TFLite model that relies on STFT processing, so that could cover many people working on projects involving audio data, financial data, or any other form of 1-dimensional data series.

Any Other info

Nothing. Just thanks to the community for keeping TensorFlow running and growing 👍

gaikwadrahul8 commented 23 hours ago

This issue originally reported by @francoisgermain has been moved to this dedicated repository for LiteRT to enhance issue tracking and prioritization. To ensure continuity, we have created this new issue on your behalf.

We appreciate your understanding and look forward to your continued involvement.