Add support for SoundStream "streaming" version
- Extend the soundstream library to let the consumer select streaming or
non-streaming variants of soundstream at construction time.
- Add build targets for streaming variants of e2e/encoder/decoder
samples (and new reference outputs, in the e2e case).
Change-Id: I45aedc11ada72390974a50f149fd08aac1ece071
diff --git a/examples/tflm/soundstream/BUILD b/examples/tflm/soundstream/BUILD
index b3a41f3..2cb2d2d 100644
--- a/examples/tflm/soundstream/BUILD
+++ b/examples/tflm/soundstream/BUILD
@@ -6,11 +6,15 @@
name = "soundstream",
srcs = [
"decoder.cc",
- "encoder.cc",
"decoder_non_stream_q16x8_b64_io_int16_tflite.cc",
- "encoder_non_stream_q16x8_b64_io_int16_tflite.cc",
"decoder_non_stream_q16x8_b64_io_int16_tflite.h",
+ "decoder_streaming_q16x8_b64_io_int16_tflite.cc",
+ "decoder_streaming_q16x8_b64_io_int16_tflite.h",
+ "encoder.cc",
+ "encoder_non_stream_q16x8_b64_io_int16_tflite.cc",
"encoder_non_stream_q16x8_b64_io_int16_tflite.h",
+ "encoder_streaming_q16x8_b64_io_int16_tflite.cc",
+ "encoder_streaming_q16x8_b64_io_int16_tflite.h",
],
hdrs = [
"decoder.h",
@@ -19,11 +23,12 @@
tags = ["manual"],
deps = [
"@tflite-micro//tensorflow/lite/micro:micro_framework",
+ "@tflite-micro//tensorflow/lite/micro:recording_allocators",
],
)
kelvin_binary(
- name = "soundstream_decoder",
+ name = "soundstream_decoder_non_streaming",
srcs = [
"soundstream_decoder.cc",
],
@@ -31,13 +36,24 @@
deps = [
":soundstream",
"//crt:crt_header",
- "@tflite-micro//tensorflow/lite/micro:micro_framework",
- "@tflite-micro//tensorflow/lite/micro:system_setup",
],
)
kelvin_binary(
- name = "soundstream_encoder",
+ name = "soundstream_decoder_streaming",
+ srcs = [
+ "soundstream_decoder.cc",
+ ],
+ copts = ["-DSTREAMING"],
+ tags = ["manual"],
+ deps = [
+ ":soundstream",
+ "//crt:crt_header",
+ ],
+)
+
+kelvin_binary(
+ name = "soundstream_encoder_non_streaming",
srcs = [
"soundstream_encoder.cc",
],
@@ -45,13 +61,24 @@
deps = [
":soundstream",
"//crt:crt_header",
- "@tflite-micro//tensorflow/lite/micro:micro_framework",
- "@tflite-micro//tensorflow/lite/micro:system_setup",
],
)
kelvin_binary(
- name = "soundstream_e2e",
+ name = "soundstream_encoder_streaming",
+ srcs = [
+ "soundstream_encoder.cc",
+ ],
+ copts = ["-DSTREAMING"],
+ tags = ["manual"],
+ deps = [
+ ":soundstream",
+ "//crt:crt_header",
+ ],
+)
+
+kelvin_binary(
+ name = "soundstream_e2e_non_streaming",
srcs = [
"best_of_times_s16_decoded.cc",
"best_of_times_s16_encoded.cc",
@@ -67,8 +94,27 @@
deps = [
":soundstream",
"//crt:crt_header",
- "@tflite-micro//tensorflow/lite/micro:micro_framework",
- "@tflite-micro//tensorflow/lite/micro:system_setup",
+ ],
+)
+
+kelvin_binary(
+ name = "soundstream_e2e_streaming",
+ srcs = [
+ "best_of_times_s16_decoded_streaming.cc",
+ "best_of_times_s16_encoded_streaming.cc",
+ "best_of_times_s16_wav.cc",
+ "soundstream_e2e.cc",
+ ],
+ hdrs = [
+ "best_of_times_s16_decoded_streaming.h",
+ "best_of_times_s16_encoded_streaming.h",
+ "best_of_times_s16_wav.h",
+ ],
+ copts = ["-DSTREAMING"],
+ tags = ["manual"],
+ deps = [
+ ":soundstream",
+ "//crt:crt_header",
],
)
@@ -101,6 +147,34 @@
)
generate_cc_arrays(
+ name = "decoder_streaming_q16x8_b64_io_int16_tflite_cc",
+ src = "@ml-models//:quant_models/_decoder_streaming_q16x8_b64_io_int16.tflite",
+ out = "decoder_streaming_q16x8_b64_io_int16_tflite.cc",
+ tags = ["manual"],
+)
+
+generate_cc_arrays(
+ name = "decoder_streaming_q16x8_b64_io_int16_tflite_h",
+ src = "@ml-models//:quant_models/_decoder_streaming_q16x8_b64_io_int16.tflite",
+ out = "decoder_streaming_q16x8_b64_io_int16_tflite.h",
+ tags = ["manual"],
+)
+
+generate_cc_arrays(
+ name = "encoder_streaming_q16x8_b64_io_int16_tflite_cc",
+ src = "@ml-models//:quant_models/_encoder_streaming_q16x8_b64_io_int16.tflite",
+ out = "encoder_streaming_q16x8_b64_io_int16_tflite.cc",
+ tags = ["manual"],
+)
+
+generate_cc_arrays(
+ name = "encoder_streaming_q16x8_b64_io_int16_tflite_h",
+ src = "@ml-models//:quant_models/_encoder_streaming_q16x8_b64_io_int16.tflite",
+ out = "encoder_streaming_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",
@@ -135,3 +209,27 @@
src = "best_of_times_s16_decoded.raw",
out = "best_of_times_s16_decoded.h",
)
+
+generate_cc_arrays(
+ name = "best_of_times_s16_encoded_streaming_cc",
+ src = "best_of_times_s16_encoded_streaming.raw",
+ out = "best_of_times_s16_encoded_streaming.cc",
+)
+
+generate_cc_arrays(
+ name = "best_of_times_s16_encoded_streaming_h",
+ src = "best_of_times_s16_encoded_streaming.raw",
+ out = "best_of_times_s16_encoded_streaming.h",
+)
+
+generate_cc_arrays(
+ name = "best_of_times_s16_decoded_streaming_cc",
+ src = "best_of_times_s16_decoded_streaming.raw",
+ out = "best_of_times_s16_decoded_streaming.cc",
+)
+
+generate_cc_arrays(
+ name = "best_of_times_s16_decoded_streaming_h",
+ src = "best_of_times_s16_decoded_streaming.raw",
+ out = "best_of_times_s16_decoded_streaming.h",
+)
diff --git a/examples/tflm/soundstream/best_of_times_s16_decoded_streaming.raw b/examples/tflm/soundstream/best_of_times_s16_decoded_streaming.raw
new file mode 100644
index 0000000..91f655e
--- /dev/null
+++ b/examples/tflm/soundstream/best_of_times_s16_decoded_streaming.raw
Binary files differ
diff --git a/examples/tflm/soundstream/best_of_times_s16_encoded_streaming.raw b/examples/tflm/soundstream/best_of_times_s16_encoded_streaming.raw
new file mode 100644
index 0000000..7956e3c
--- /dev/null
+++ b/examples/tflm/soundstream/best_of_times_s16_encoded_streaming.raw
Binary files differ
diff --git a/examples/tflm/soundstream/decoder.cc b/examples/tflm/soundstream/decoder.cc
index 2a0fc4c..6aa939a 100644
--- a/examples/tflm/soundstream/decoder.cc
+++ b/examples/tflm/soundstream/decoder.cc
@@ -1,37 +1,96 @@
+// 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/decoder.h"
#include "examples/tflm/soundstream/decoder_non_stream_q16x8_b64_io_int16_tflite.h"
+#include "examples/tflm/soundstream/decoder_streaming_q16x8_b64_io_int16_tflite.h"
namespace kelvin::soundstream::decoder {
-std::optional<Decoder> Setup(uint8_t* tensor_arena) {
- auto* model =
- tflite::GetModel(g__decoder_non_stream_q16x8_b64_io_int16_model_data);
- if (model->version() != TFLITE_SCHEMA_VERSION) {
- return {};
+
+constexpr unsigned int kNonStreamingOpCount = 11;
+constexpr unsigned int kStreamingOpCount = 16;
+// Not sure how to get a good upper bound on this one, so arbitrarily chosen.
+constexpr unsigned int kStreamingVariablesCount = 40;
+
+template <bool kStreaming>
+class DecoderImpl : public Decoder {
+ public:
+ static Decoder* Setup(const uint8_t* model_data, uint8_t* tensor_arena,
+ size_t tensor_arena_size) {
+ auto* model = tflite::GetModel(model_data);
+ if (model->version() != TFLITE_SCHEMA_VERSION) {
+ return nullptr;
+ }
+
+ DecoderImpl* d = new DecoderImpl(model, tensor_arena, tensor_arena_size);
+
+ TfLiteStatus allocate_status = d->interpreter()->AllocateTensors();
+ if (allocate_status != kTfLiteOk) {
+ MicroPrintf("Failed to allocate decoder's tensors");
+ return nullptr;
+ }
+ return d;
}
+ tflite::MicroInterpreter* interpreter() { return &interpreter_; }
- Decoder d;
- d.resolver = std::make_unique<tflite::MicroMutableOpResolver<11>>();
- d.resolver->AddReshape();
- d.resolver->AddPad();
- d.resolver->AddConv2D();
- d.resolver->AddLeakyRelu();
- d.resolver->AddSplit();
- d.resolver->AddTransposeConv();
- d.resolver->AddStridedSlice();
- d.resolver->AddConcatenation();
- d.resolver->AddDepthwiseConv2D();
- d.resolver->AddAdd();
- d.resolver->AddQuantize();
+ private:
+ DecoderImpl(const tflite::Model* model, uint8_t* tensor_arena,
+ size_t tensor_arena_size)
+ : resolver_(CreateResolver()),
+ allocator_(tflite::RecordingMicroAllocator::Create(tensor_arena,
+ tensor_arena_size)),
+ variables_(tflite::MicroResourceVariables::Create(
+ allocator_.get(), kStreamingVariablesCount)),
+ interpreter_(model, resolver_, allocator_.get(), variables_.get()) {}
- d.interpreter = std::make_unique<tflite::MicroInterpreter>(
- model, *d.resolver, tensor_arena, kTensorArenaSizeBytes);
-
- TfLiteStatus allocate_status = d.interpreter->AllocateTensors();
- if (allocate_status != kTfLiteOk) {
- MicroPrintf("Failed to allocate decoder's tensors");
- return {};
+ static constexpr int kOpCount =
+ kStreaming ? kStreamingOpCount : kStreamingOpCount;
+ static inline tflite::MicroMutableOpResolver<kOpCount> CreateResolver() {
+ tflite::MicroMutableOpResolver<kOpCount> resolver;
+ resolver.AddReshape();
+ resolver.AddPad();
+ resolver.AddConv2D();
+ resolver.AddLeakyRelu();
+ resolver.AddSplit();
+ resolver.AddTransposeConv();
+ resolver.AddStridedSlice();
+ resolver.AddConcatenation();
+ resolver.AddDepthwiseConv2D();
+ resolver.AddAdd();
+ resolver.AddQuantize();
+ if (kStreaming) {
+ resolver.AddCallOnce();
+ resolver.AddVarHandle();
+ resolver.AddReadVariable();
+ resolver.AddAssignVariable();
+ resolver.AddSub();
+ }
+ return resolver;
}
- return d;
+ const tflite::MicroMutableOpResolver<kOpCount> resolver_;
+ // Created in the arena
+ std::unique_ptr<tflite::RecordingMicroAllocator> allocator_;
+ // Created in the arena
+ std::unique_ptr<tflite::MicroResourceVariables> variables_;
+ tflite::MicroInterpreter interpreter_;
+};
+
+// Two separate methods to construct streaming vs non-streaming, so that the
+// compiler can eliminate one if it's unused. Perhaps with LTO we could combine
+// them together.
+std::unique_ptr<Decoder> SetupStreaming(uint8_t* tensor_arena,
+ size_t tensor_arena_size) {
+ return std::unique_ptr<Decoder>(DecoderImpl<true>::Setup(
+ g__decoder_streaming_q16x8_b64_io_int16_model_data, tensor_arena,
+ tensor_arena_size));
}
+std::unique_ptr<Decoder> Setup(uint8_t* tensor_arena,
+ size_t tensor_arena_size) {
+ return std::unique_ptr<Decoder>(DecoderImpl<false>::Setup(
+ g__decoder_non_stream_q16x8_b64_io_int16_model_data, tensor_arena,
+ tensor_arena_size));
+}
+
} // namespace kelvin::soundstream::decoder
diff --git a/examples/tflm/soundstream/decoder.h b/examples/tflm/soundstream/decoder.h
index b21c93d..337f402 100644
--- a/examples/tflm/soundstream/decoder.h
+++ b/examples/tflm/soundstream/decoder.h
@@ -1,20 +1,31 @@
+// Copyright 2023 Google LLC
+// Licensed under the Apache License, Version 2.0, see LICENSE for details.
+// SPDX-License-Identifier: Apache-2.0
+
#ifndef EXAMPLES_TFLM_SOUNDSTREAM_DECODER_H_
#define EXAMPLES_TFLM_SOUNDSTREAM_DECODER_H_
#include <cstddef>
#include <memory>
-#include <optional>
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/micro/micro_mutable_op_resolver.h"
+#include "tensorflow/lite/micro/recording_micro_allocator.h"
namespace kelvin::soundstream::decoder {
+// RecordingMicroAllocator on desktop recorded 94512 bytes of allocation.
constexpr size_t kTensorArenaSizeBytes = 96 * 1024;
-struct Decoder {
- std::unique_ptr<tflite::MicroInterpreter> interpreter;
- std::unique_ptr<tflite::MicroMutableOpResolver<11>> resolver;
+// RecordingMicroAllocator on desktop recorded 143296 bytes of allocation.
+constexpr size_t kTensorArenaStreamingSizeBytes = 168 * 1024;
+
+class Decoder {
+ public:
+ virtual tflite::MicroInterpreter* interpreter() = 0;
};
-std::optional<Decoder> Setup(uint8_t* tensor_arena);
+
+std::unique_ptr<Decoder> Setup(uint8_t* tensor_arena, size_t tensor_arena_size);
+std::unique_ptr<Decoder> SetupStreaming(uint8_t* tensor_arena,
+ size_t tensor_arena_size);
} // namespace kelvin::soundstream::decoder
#endif // EXAMPLES_TFLM_SOUNDSTREAM_DECODER_H_
diff --git a/examples/tflm/soundstream/encoder.cc b/examples/tflm/soundstream/encoder.cc
index d31d517..6d3c985 100644
--- a/examples/tflm/soundstream/encoder.cc
+++ b/examples/tflm/soundstream/encoder.cc
@@ -1,32 +1,95 @@
+// 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/encoder.h"
#include "examples/tflm/soundstream/encoder_non_stream_q16x8_b64_io_int16_tflite.h"
+#include "examples/tflm/soundstream/encoder_streaming_q16x8_b64_io_int16_tflite.h"
+#include "tensorflow/lite/micro/micro_mutable_op_resolver.h"
+#include "tensorflow/lite/micro/recording_micro_allocator.h"
namespace kelvin::soundstream::encoder {
-std::optional<Encoder> Setup(uint8_t* tensor_arena) {
- auto* model =
- tflite::GetModel(g__encoder_non_stream_q16x8_b64_io_int16_model_data);
- if (model->version() != TFLITE_SCHEMA_VERSION) {
- return {};
+
+constexpr unsigned int kNonStreamingOpCount = 6;
+constexpr unsigned int kStreamingOpCount = 13;
+// Not sure how to get a good upper bound on this one, so arbitrarily chosen.
+constexpr unsigned int kStreamingVariablesCount = 40;
+
+template <bool kStreaming>
+class EncoderImpl : public Encoder {
+ public:
+ static Encoder* Setup(const uint8_t* model_data, uint8_t* tensor_arena,
+ size_t tensor_arena_size) {
+ auto* model = tflite::GetModel(model_data);
+ if (model->version() != TFLITE_SCHEMA_VERSION) {
+ return nullptr;
+ }
+
+ EncoderImpl* e = new EncoderImpl(model, tensor_arena, tensor_arena_size);
+
+ TfLiteStatus allocate_status = e->interpreter()->AllocateTensors();
+ if (allocate_status != kTfLiteOk) {
+ MicroPrintf("Failed to allocate decoder's tensors");
+ return nullptr;
+ }
+ return e;
}
+ tflite::MicroInterpreter* interpreter() { return &interpreter_; }
- Encoder e;
- e.resolver = std::make_unique<tflite::MicroMutableOpResolver<6>>();
- e.resolver->AddReshape();
- e.resolver->AddPad();
- e.resolver->AddConv2D();
- e.resolver->AddLeakyRelu();
- e.resolver->AddDepthwiseConv2D();
- e.resolver->AddAdd();
+ private:
+ EncoderImpl(const tflite::Model* model, uint8_t* tensor_arena,
+ size_t tensor_arena_size)
+ : resolver_(CreateResolver()),
+ allocator_(tflite::RecordingMicroAllocator::Create(tensor_arena,
+ tensor_arena_size)),
+ variables_(tflite::MicroResourceVariables::Create(
+ allocator_.get(), kStreamingVariablesCount)),
+ interpreter_(model, resolver_, allocator_.get(), variables_.get()) {}
- e.interpreter = std::make_unique<tflite::MicroInterpreter>(
- model, *e.resolver, tensor_arena, kTensorArenaSizeBytes);
-
- TfLiteStatus allocate_status = e.interpreter->AllocateTensors();
- if (allocate_status != kTfLiteOk) {
- MicroPrintf("Failed to allocate encoder's tensors");
- return {};
+ static constexpr int kOpCount =
+ kStreaming ? kStreamingOpCount : kStreamingOpCount;
+ static inline tflite::MicroMutableOpResolver<kOpCount> CreateResolver() {
+ tflite::MicroMutableOpResolver<kOpCount> resolver;
+ resolver.AddReshape();
+ resolver.AddPad();
+ resolver.AddConv2D();
+ resolver.AddLeakyRelu();
+ resolver.AddDepthwiseConv2D();
+ resolver.AddAdd();
+ if (kStreaming) {
+ resolver.AddCallOnce();
+ resolver.AddVarHandle();
+ resolver.AddReadVariable();
+ resolver.AddConcatenation();
+ resolver.AddStridedSlice();
+ resolver.AddAssignVariable();
+ resolver.AddQuantize();
+ }
+ return resolver;
}
- return e;
+ const tflite::MicroMutableOpResolver<kOpCount> resolver_;
+ // Created in the arena
+ std::unique_ptr<tflite::RecordingMicroAllocator> allocator_;
+ // Created in the arena
+ std::unique_ptr<tflite::MicroResourceVariables> variables_;
+ tflite::MicroInterpreter interpreter_;
+};
+
+// Two separate methods to construct streaming vs non-streaming, so that the
+// compiler can eliminate one if it's unused. Perhaps with LTO we could combine
+// them together.
+std::unique_ptr<Encoder> SetupStreaming(uint8_t* tensor_arena,
+ size_t tensor_arena_size) {
+ return std::unique_ptr<Encoder>(EncoderImpl<true>::Setup(
+ g__encoder_streaming_q16x8_b64_io_int16_model_data, tensor_arena,
+ tensor_arena_size));
}
+std::unique_ptr<Encoder> Setup(uint8_t* tensor_arena,
+ size_t tensor_arena_size) {
+ return std::unique_ptr<Encoder>(EncoderImpl<false>::Setup(
+ g__encoder_non_stream_q16x8_b64_io_int16_model_data, tensor_arena,
+ tensor_arena_size));
+}
+
} // namespace kelvin::soundstream::encoder
diff --git a/examples/tflm/soundstream/encoder.h b/examples/tflm/soundstream/encoder.h
index 3cb74f9..267d553 100644
--- a/examples/tflm/soundstream/encoder.h
+++ b/examples/tflm/soundstream/encoder.h
@@ -1,20 +1,28 @@
+// Copyright 2023 Google LLC
+// Licensed under the Apache License, Version 2.0, see LICENSE for details.
+// SPDX-License-Identifier: Apache-2.0
+
#ifndef EXAMPLES_TFLM_SOUNDSTREAM_ENCODER_H_
#define EXAMPLES_TFLM_SOUNDSTREAM_ENCODER_H_
#include <cstddef>
#include <memory>
-#include <optional>
#include "tensorflow/lite/micro/micro_interpreter.h"
-#include "tensorflow/lite/micro/micro_mutable_op_resolver.h"
namespace kelvin::soundstream::encoder {
+// RecordingMicroAllocator on desktop recorded 90064 bytes of allocation.
constexpr size_t kTensorArenaSizeBytes = 96 * 1024;
+// RecordingMicroAllocator on desktop recorded 147328 bytes of allocation.
+constexpr size_t kTensorArenaStreamingSizeBytes = 168 * 1024;
+
struct Encoder {
- std::unique_ptr<tflite::MicroInterpreter> interpreter;
- std::unique_ptr<tflite::MicroMutableOpResolver<6>> resolver;
+ virtual tflite::MicroInterpreter* interpreter() = 0;
};
-std::optional<Encoder> Setup(uint8_t* tensor_arena);
+
+std::unique_ptr<Encoder> Setup(uint8_t* tensor_arena, size_t tensor_arena_size);
+std::unique_ptr<Encoder> SetupStreaming(uint8_t* tensor_arena,
+ size_t tensor_arena_size);
} // namespace kelvin::soundstream::encoder
#endif // EXAMPLES_TFLM_SOUNDSTREAM_ENCODER_H_
diff --git a/examples/tflm/soundstream/soundstream_decoder.cc b/examples/tflm/soundstream/soundstream_decoder.cc
index 77f0f94..9a5a739 100644
--- a/examples/tflm/soundstream/soundstream_decoder.cc
+++ b/examples/tflm/soundstream/soundstream_decoder.cc
@@ -16,23 +16,34 @@
};
namespace {
-uint8_t
- decoder_tensor_arena[kelvin::soundstream::decoder::kTensorArenaSizeBytes]
- __attribute__((aligned(64)));
+#if defined(STREAMING)
+constexpr size_t tensor_arena_size =
+ kelvin::soundstream::decoder::kTensorArenaStreamingSizeBytes;
+#else
+constexpr size_t tensor_arena_size =
+ kelvin::soundstream::decoder::kTensorArenaSizeBytes;
+#endif
+uint8_t decoder_tensor_arena[tensor_arena_size] __attribute__((aligned(64)));
} // namespace
int main(int argc, char **argv) {
- auto decoder = kelvin::soundstream::decoder::Setup(decoder_tensor_arena);
+#if defined(STREAMING)
+ auto decoder = kelvin::soundstream::decoder::SetupStreaming(
+ decoder_tensor_arena, tensor_arena_size);
+#else
+ auto decoder = kelvin::soundstream::decoder::Setup(decoder_tensor_arena,
+ tensor_arena_size);
+#endif
if (!decoder) {
MicroPrintf("Unable to construct decoder");
return -1;
}
- TfLiteTensor *decoder_input = decoder->interpreter->input(0);
- TfLiteTensor *decoder_output = decoder->interpreter->output(0);
+ TfLiteTensor *decoder_input = decoder->interpreter()->input(0);
+ TfLiteTensor *decoder_output = decoder->interpreter()->output(0);
memset(decoder_input->data.uint8, 0, decoder_input->bytes);
- TfLiteStatus invoke_status = decoder->interpreter->Invoke();
+ TfLiteStatus invoke_status = decoder->interpreter()->Invoke();
if (invoke_status != kTfLiteOk) {
MicroPrintf("Failed to invoke decoder");
return -1;
diff --git a/examples/tflm/soundstream/soundstream_e2e.cc b/examples/tflm/soundstream/soundstream_e2e.cc
index 4c43060..9ba8913 100644
--- a/examples/tflm/soundstream/soundstream_e2e.cc
+++ b/examples/tflm/soundstream/soundstream_e2e.cc
@@ -2,38 +2,69 @@
// 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_decoded.h"
-#include "examples/tflm/soundstream/best_of_times_s16_encoded.h"
#include "examples/tflm/soundstream/best_of_times_s16_wav.h"
#include "examples/tflm/soundstream/decoder.h"
#include "examples/tflm/soundstream/encoder.h"
+#if defined(STREAMING)
+#include "examples/tflm/soundstream/best_of_times_s16_decoded_streaming.h"
+#include "examples/tflm/soundstream/best_of_times_s16_encoded_streaming.h"
+const unsigned char *reference_decoded = g_best_of_times_s16_decoded_streaming;
+const unsigned char *reference_encoded = g_best_of_times_s16_encoded_streaming;
+#else
+#include "examples/tflm/soundstream/best_of_times_s16_decoded.h"
+#include "examples/tflm/soundstream/best_of_times_s16_encoded.h"
+const unsigned char *reference_decoded = g_best_of_times_s16_decoded;
+const unsigned char *reference_encoded = g_best_of_times_s16_encoded;
+#endif
+
namespace {
-uint8_t
- encoder_tensor_arena[kelvin::soundstream::encoder::kTensorArenaSizeBytes]
+#if defined(STREAMING)
+constexpr size_t decoder_tensor_arena_size =
+ kelvin::soundstream::decoder::kTensorArenaStreamingSizeBytes;
+constexpr size_t encoder_tensor_arena_size =
+ kelvin::soundstream::encoder::kTensorArenaStreamingSizeBytes;
+#else
+constexpr size_t decoder_tensor_arena_size =
+ kelvin::soundstream::decoder::kTensorArenaSizeBytes;
+constexpr size_t encoder_tensor_arena_size =
+ kelvin::soundstream::encoder::kTensorArenaSizeBytes;
+#endif
+uint8_t encoder_tensor_arena[encoder_tensor_arena_size]
__attribute__((aligned(64)));
-uint8_t
- decoder_tensor_arena[kelvin::soundstream::decoder::kTensorArenaSizeBytes]
+uint8_t decoder_tensor_arena[decoder_tensor_arena_size]
__attribute__((aligned(64)));
} // namespace
int main(int argc, char **argv) {
- auto encoder = kelvin::soundstream::encoder::Setup(encoder_tensor_arena);
+#if defined(STREAMING)
+ auto encoder = kelvin::soundstream::encoder::SetupStreaming(
+ encoder_tensor_arena, encoder_tensor_arena_size);
+#else
+ auto encoder = kelvin::soundstream::encoder::Setup(encoder_tensor_arena,
+ encoder_tensor_arena_size);
+#endif
if (!encoder) {
MicroPrintf("Unable to construct encoder");
return -1;
}
- auto decoder = kelvin::soundstream::decoder::Setup(decoder_tensor_arena);
+#if defined(STREAMING)
+ auto decoder = kelvin::soundstream::decoder::SetupStreaming(
+ decoder_tensor_arena, decoder_tensor_arena_size);
+#else
+ auto decoder = kelvin::soundstream::decoder::Setup(decoder_tensor_arena,
+ decoder_tensor_arena_size);
+#endif
if (!decoder) {
MicroPrintf("Unable to construct decoder");
return -1;
}
- TfLiteTensor *encoder_input = encoder->interpreter->input(0);
- TfLiteTensor *encoder_output = encoder->interpreter->output(0);
- TfLiteTensor *decoder_input = decoder->interpreter->input(0);
- TfLiteTensor *decoder_output = decoder->interpreter->output(0);
+ TfLiteTensor *encoder_input = encoder->interpreter()->input(0);
+ TfLiteTensor *encoder_output = encoder->interpreter()->output(0);
+ TfLiteTensor *decoder_input = decoder->interpreter()->input(0);
+ TfLiteTensor *decoder_output = decoder->interpreter()->output(0);
int invocation_count =
(g_best_of_times_s16_audio_data_size * sizeof(int16_t)) /
@@ -44,13 +75,13 @@
g_best_of_times_s16_audio_data +
((i * encoder_input->bytes) / sizeof(int16_t)),
encoder_input->bytes);
- TfLiteStatus invoke_status = encoder->interpreter->Invoke();
+ TfLiteStatus invoke_status = encoder->interpreter()->Invoke();
if (invoke_status != kTfLiteOk) {
MicroPrintf("Failed to invoke encoder");
return -1;
}
if (memcmp(encoder_output->data.uint8,
- g_best_of_times_s16_encoded + (i * encoder_output->bytes),
+ reference_encoded + (i * encoder_output->bytes),
encoder_output->bytes)) {
MicroPrintf("Encoder output mismatches reference");
return -1;
@@ -58,13 +89,13 @@
memcpy(decoder_input->data.uint8, encoder_output->data.uint8,
decoder_input->bytes);
- invoke_status = decoder->interpreter->Invoke();
+ invoke_status = decoder->interpreter()->Invoke();
if (invoke_status != kTfLiteOk) {
MicroPrintf("Failed to invoke decoder");
return -1;
}
if (memcmp(decoder_output->data.uint8,
- g_best_of_times_s16_decoded + (i * decoder_output->bytes),
+ reference_decoded + (i * decoder_output->bytes),
decoder_output->bytes)) {
MicroPrintf("Decoder output mismatches reference");
return -1;
diff --git a/examples/tflm/soundstream/soundstream_encoder.cc b/examples/tflm/soundstream/soundstream_encoder.cc
index bc7fc25..6c60748 100644
--- a/examples/tflm/soundstream/soundstream_encoder.cc
+++ b/examples/tflm/soundstream/soundstream_encoder.cc
@@ -16,23 +16,34 @@
};
namespace {
-uint8_t
- encoder_tensor_arena[kelvin::soundstream::encoder::kTensorArenaSizeBytes]
- __attribute__((aligned(64)));
+#if defined(STREAMING)
+constexpr size_t tensor_arena_size =
+ kelvin::soundstream::encoder::kTensorArenaStreamingSizeBytes;
+#else
+constexpr size_t tensor_arena_size =
+ kelvin::soundstream::encoder::kTensorArenaSizeBytes;
+#endif
+uint8_t tensor_arena[tensor_arena_size] __attribute__((aligned(64)));
} // namespace
int main(int argc, char **argv) {
- auto encoder = kelvin::soundstream::encoder::Setup(encoder_tensor_arena);
+#if defined(STREAMING)
+ auto encoder = kelvin::soundstream::encoder::SetupStreaming(
+ tensor_arena, tensor_arena_size);
+#else
+ auto encoder =
+ kelvin::soundstream::encoder::Setup(tensor_arena, tensor_arena_size);
+#endif
if (!encoder) {
MicroPrintf("Unable to construct encoder");
return -1;
}
- TfLiteTensor *encoder_input = encoder->interpreter->input(0);
- TfLiteTensor *encoder_output = encoder->interpreter->output(0);
+ TfLiteTensor *encoder_input = encoder->interpreter()->input(0);
+ TfLiteTensor *encoder_output = encoder->interpreter()->output(0);
memset(encoder_input->data.uint8, 0, encoder_input->bytes);
- TfLiteStatus invoke_status = encoder->interpreter->Invoke();
+ TfLiteStatus invoke_status = encoder->interpreter()->Invoke();
if (invoke_status != kTfLiteOk) {
MicroPrintf("Failed to invoke encoder");
return -1;
