Soundstream encode/decode in TFLM - Add a sample application that runs a fixed input through the soundstream encoder, and then runs that output back through the decoder. - Import a patch for TFLM, which adds depthwise conv support for more datatypes. Change-Id: I7fad53595b105e001963f9d4d1c9106e2cd84344
diff --git a/WORKSPACE b/WORKSPACE index 5f7eedf..8a677b0 100644 --- a/WORKSPACE +++ b/WORKSPACE
@@ -1,6 +1,6 @@ workspace(name = "kelvin_sw") -load("//build_tools/bazel:repos.bzl", "kelvin_repos", "tflm_repos") +load("//build_tools/bazel:repos.bzl", "kelvin_repos", "model_repos", "tflm_repos") kelvin_repos() @@ -22,3 +22,5 @@ load("@tflm_pip_deps//:requirements.bzl", "install_deps") install_deps() + +model_repos()
diff --git a/build_tools/bazel/repos.bzl b/build_tools/bazel/repos.bzl index 5cc5b76..20ba4b8 100644 --- a/build_tools/bazel/repos.bzl +++ b/build_tools/bazel/repos.bzl
@@ -90,3 +90,10 @@ urls = ["https://github.com/pybind/pybind11_bazel/archive/faf56fb3df11287f26dbc66fdedf60a2fc2c6631.zip"], sha256 = "a185aa68c93b9f62c80fcb3aadc3c83c763854750dc3f38be1dadcb7be223837", ) + +def model_repos(): + native.new_local_repository( + name = "ml-models", + path = "../../ml/ml-models", + build_file = "third_party/ml-models/BUILD", + )
diff --git a/examples/tflm/soundstream/BUILD b/examples/tflm/soundstream/BUILD new file mode 100644 index 0000000..1188d98 --- /dev/null +++ b/examples/tflm/soundstream/BUILD
@@ -0,0 +1,60 @@ +load("//build_tools/bazel:kelvin.bzl", "generate_cc_arrays", "kelvin_binary") +package(default_visibility = ["//visibility:public"]) + +kelvin_binary( + name = "soundstream", + srcs = [ + "soundstream.cc", + "best_of_times_s16_wav.cc", + "decoder_non_stream_q16x8_b64_io_int16_tflite.cc", + "encoder_non_stream_q16x8_b64_io_int16_tflite.cc", + ], + hdrs = [ + "best_of_times_s16_wav.h", + "decoder_non_stream_q16x8_b64_io_int16_tflite.h", + "encoder_non_stream_q16x8_b64_io_int16_tflite.h", + ], + deps = [ + "//crt:crt_header", + "@tflite-micro//tensorflow/lite/micro:micro_framework", + "@tflite-micro//tensorflow/lite/micro:system_setup", + ], + tags = ["manual"], +) + +generate_cc_arrays( + name = "decoder_non_stream_q16x8_b64_io_int16_tflite_cc", + src = "@ml-models//:quant_models/_decoder_non_stream_q16x8_b64_io_int16.tflite", + out = "decoder_non_stream_q16x8_b64_io_int16_tflite.cc", + tags = ["manual"], +) +generate_cc_arrays( + name = "decoder_non_stream_q16x8_b64_io_int16_tflite_h", + src = "@ml-models//:quant_models/_decoder_non_stream_q16x8_b64_io_int16.tflite", + out = "decoder_non_stream_q16x8_b64_io_int16_tflite.h", + tags = ["manual"], +) + +generate_cc_arrays( + name = "encoder_non_stream_q16x8_b64_io_int16_tflite_cc", + src = "@ml-models//:quant_models/_encoder_non_stream_q16x8_b64_io_int16.tflite", + out = "encoder_non_stream_q16x8_b64_io_int16_tflite.cc", + tags = ["manual"], +) +generate_cc_arrays( + name = "encoder_non_stream_q16x8_b64_io_int16_tflite_h", + src = "@ml-models//:quant_models/_encoder_non_stream_q16x8_b64_io_int16.tflite", + out = "encoder_non_stream_q16x8_b64_io_int16_tflite.h", + tags = ["manual"], +) + +generate_cc_arrays( + name = "best_of_times_s16_wav_cc", + src = "best_of_times_s16.wav", + out = "best_of_times_s16_wav.cc", +) +generate_cc_arrays( + name = "best_of_times_s16_wav_h", + src = "best_of_times_s16.wav", + out = "best_of_times_s16_wav.h", +)
diff --git a/examples/tflm/soundstream/best_of_times_s16.wav b/examples/tflm/soundstream/best_of_times_s16.wav new file mode 100644 index 0000000..e3a7685 --- /dev/null +++ b/examples/tflm/soundstream/best_of_times_s16.wav Binary files differ
diff --git a/examples/tflm/soundstream/soundstream.cc b/examples/tflm/soundstream/soundstream.cc new file mode 100644 index 0000000..b2327fd --- /dev/null +++ b/examples/tflm/soundstream/soundstream.cc
@@ -0,0 +1,102 @@ +// Copyright 2023 Google LLC +// Licensed under the Apache License, Version 2.0, see LICENSE for details. +// SPDX-License-Identifier: Apache-2.0 + +#include "examples/tflm/soundstream/best_of_times_s16_wav.h" +#include "examples/tflm/soundstream/decoder_non_stream_q16x8_b64_io_int16_tflite.h" +#include "examples/tflm/soundstream/encoder_non_stream_q16x8_b64_io_int16_tflite.h" +#include "tensorflow/lite/micro/micro_interpreter.h" +#include "tensorflow/lite/micro/micro_mutable_op_resolver.h" + +namespace { +const tflite::Model *encoder_model = nullptr; +const tflite::Model *decoder_model = nullptr; +tflite::MicroInterpreter *encoder_interpreter = nullptr; +tflite::MicroInterpreter *decoder_interpreter = nullptr; +constexpr int kTensorArenaSize = + 96 * 1024; +uint8_t encoder_tensor_arena[kTensorArenaSize] __attribute__((aligned(16))); +uint8_t decoder_tensor_arena[kTensorArenaSize] __attribute__((aligned(16))); +} // namespace + +int main(int argc, char **argv) { + encoder_model = + tflite::GetModel(g__encoder_non_stream_q16x8_b64_io_int16_model_data); + if (encoder_model->version() != TFLITE_SCHEMA_VERSION) { + return 1; + } + decoder_model = + tflite::GetModel(g__decoder_non_stream_q16x8_b64_io_int16_model_data); + if (decoder_model->version() != TFLITE_SCHEMA_VERSION) { + return 1; + } + + static tflite::MicroMutableOpResolver<6> encoder_resolver{}; + encoder_resolver.AddReshape(); + encoder_resolver.AddPad(); + encoder_resolver.AddConv2D(); + encoder_resolver.AddLeakyRelu(); + encoder_resolver.AddDepthwiseConv2D(); + encoder_resolver.AddAdd(); + + static tflite::MicroMutableOpResolver<11> decoder_resolver{}; + decoder_resolver.AddReshape(); + decoder_resolver.AddPad(); + decoder_resolver.AddConv2D(); + decoder_resolver.AddLeakyRelu(); + decoder_resolver.AddSplit(); + decoder_resolver.AddTransposeConv(); + decoder_resolver.AddStridedSlice(); + decoder_resolver.AddConcatenation(); + decoder_resolver.AddDepthwiseConv2D(); + decoder_resolver.AddAdd(); + decoder_resolver.AddQuantize(); + + static tflite::MicroInterpreter encoder_static_interpreter( + encoder_model, encoder_resolver, encoder_tensor_arena, kTensorArenaSize); + encoder_interpreter = &encoder_static_interpreter; + + static tflite::MicroInterpreter decoder_static_interpreter( + decoder_model, decoder_resolver, decoder_tensor_arena, kTensorArenaSize); + decoder_interpreter = &decoder_static_interpreter; + + TfLiteStatus allocate_status = encoder_interpreter->AllocateTensors(); + if (allocate_status != kTfLiteOk) { + MicroPrintf("Failed to allocate encoder's tensors"); + return -1; + } + allocate_status = decoder_interpreter->AllocateTensors(); + if (allocate_status != kTfLiteOk) { + MicroPrintf("Failed to allocate decoder's tensors"); + return -1; + } + + TfLiteTensor *encoder_input = encoder_interpreter->input(0); + TfLiteTensor *encoder_output = encoder_interpreter->output(0); + + int invocation_count = + g_best_of_times_s16_audio_data_size / encoder_input->bytes; + for (int i = 0; i < invocation_count; ++i) { + MicroPrintf("Invocation %d of %d", i, invocation_count); + memcpy(encoder_input->data.uint8, + g_best_of_times_s16_audio_data + + ((i * encoder_input->bytes) / sizeof(int16_t)), + encoder_input->bytes); + TfLiteStatus invoke_status = encoder_interpreter->Invoke(); + if (invoke_status != kTfLiteOk) { + MicroPrintf("Failed to invoke encoder"); + return -1; + } + + TfLiteTensor *decoder_input = decoder_interpreter->input(0); + memcpy(decoder_input->data.uint8, encoder_output->data.uint8, + decoder_input->bytes); + invoke_status = decoder_interpreter->Invoke(); + if (invoke_status != kTfLiteOk) { + MicroPrintf("Failed to invoke decoder"); + return -1; + } + } + + return 0; +}
diff --git a/third_party/ml-models/BUILD b/third_party/ml-models/BUILD new file mode 100644 index 0000000..3f152f8 --- /dev/null +++ b/third_party/ml-models/BUILD
@@ -0,0 +1,6 @@ +# Copyright 2023 Google LLC +package(default_visibility = ["//visibility:public"]) + +exports_files( + glob(["quant_models/*.tflite"]) +)