From 8a95b0e17024a182af75082dcfd1f7946b15b8ab Mon Sep 17 00:00:00 2001 From: wj <126436871+WenJing95@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:35:47 +0900 Subject: [PATCH] fix(whisper): match mel frontend boundaries Signed-off-by: wj <126436871+WenJing95@users.noreply.github.com> --- families/whisper/runtime/CMakeLists.txt | 8 ++ .../runtime/whisper_mel_spectrogram.cpp | 8 +- .../cpp/test_whisper_mel_spectrogram.cpp | 127 ++++++++++++++++++ 3 files changed, 141 insertions(+), 2 deletions(-) create mode 100644 families/whisper/tests/cpp/test_whisper_mel_spectrogram.cpp diff --git a/families/whisper/runtime/CMakeLists.txt b/families/whisper/runtime/CMakeLists.txt index c14d87740e..51f108cfbe 100644 --- a/families/whisper/runtime/CMakeLists.txt +++ b/families/whisper/runtime/CMakeLists.txt @@ -41,6 +41,14 @@ install(TARGETS trtmc_model_whisper ) if(TRTMC_BUILD_TESTS) + add_executable(test_whisper_mel_spectrogram + ${PROJECT_SOURCE_DIR}/families/whisper/tests/cpp/test_whisper_mel_spectrogram.cpp + whisper_mel_spectrogram.cpp + ) + target_include_directories(test_whisper_mel_spectrogram PRIVATE ${PROJECT_SOURCE_DIR}) + target_compile_options(test_whisper_mel_spectrogram PRIVATE -Wall -Wextra -Wpedantic) + add_test(NAME whisper_mel_spectrogram COMMAND test_whisper_mel_spectrogram) + # CPU CUDA stubs live in the test, so allocation failures can be injected # without a GPU or cudart. add_executable(test_whisper_cross_kv_alloc diff --git a/families/whisper/runtime/whisper_mel_spectrogram.cpp b/families/whisper/runtime/whisper_mel_spectrogram.cpp index 69e023128e..36c4d6177e 100644 --- a/families/whisper/runtime/whisper_mel_spectrogram.cpp +++ b/families/whisper/runtime/whisper_mel_spectrogram.cpp @@ -161,6 +161,11 @@ std::vector build_center_padded_audio(const float* samples, int32_t n_sam const int32_t padded_length = pad_size + audio_length + pad_size; std::vector padded(padded_length, 0.0F); std::memcpy(padded.data() + pad_size, audio_padded.data(), audio_length * sizeof(float)); + // STFT centers the fixed-length waveform with reflection, excluding each edge sample. + for (int32_t i = 0; i < pad_size; ++i) { + padded[i] = audio_padded[pad_size - i]; + padded[pad_size + audio_length + i] = audio_padded[audio_length - 2 - i]; + } return padded; } @@ -274,10 +279,9 @@ MelResult extract_mel_spectrogram(const float* samples, int32_t n_samples, const std::vector mel_spec = compute_mel_spectrogram(padded, make_hann_window(n_fft), mel_filters, n_fft, hop_length, freq_bins, n_mel_bins, frames_to_compute, n_frames_raw); - normalize_log_mel_inplace(mel_spec); - int32_t n_frames_out = 0; mel_spec = trim_last_frame(std::move(mel_spec), n_mel_bins, n_frames_raw, n_frames_out); + normalize_log_mel_inplace(mel_spec); MelResult result; result.data = std::move(mel_spec); diff --git a/families/whisper/tests/cpp/test_whisper_mel_spectrogram.cpp b/families/whisper/tests/cpp/test_whisper_mel_spectrogram.cpp new file mode 100644 index 0000000000..003a4af810 --- /dev/null +++ b/families/whisper/tests/cpp/test_whisper_mel_spectrogram.cpp @@ -0,0 +1,127 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/whisper/runtime/whisper_mel_spectrogram.h" + +#include +#include +#include +#include +#include + +namespace { + +int failures = 0; + +void check(bool condition, const char* message) { + if (!condition) { + std::fprintf(stderr, "FAIL: %s\n", message); + ++failures; + } +} + +struct Golden { + int mel; + int frame; + float value; +}; + +void check_values(const trtmc::whisper::MelResult& result, std::initializer_list expected) { + const bool valid_shape = (result.n_mels == 80 || result.n_mels == 128) && + result.n_frames == 3000 && + result.data.size() == static_cast(result.n_mels) * 3000; + check(valid_shape, "30-second output has the expected mel shape"); + if (!valid_shape) { + return; + } + for (const auto& point : expected) { + const float actual = result.data[point.mel * result.n_frames + point.frame]; + if (!std::isfinite(actual) || std::abs(actual - point.value) > 3e-5F) { + std::fprintf(stderr, "FAIL: mel %d frame %d: expected %.8f, got %.8f\n", point.mel, + point.frame, point.value, actual); + ++failures; + } + } +} + +trtmc::whisper::MelResult extract(const std::vector& audio, int mels) { + // A sparse test filterbank selects FFT bin m + 1 for mel m. + std::vector filters(201 * mels, 0.0F); + for (int m = 0; m < mels; ++m) { + filters[(m + 1) * mels + m] = 1.0F; + } + return trtmc::whisper::extract_mel_spectrogram(audio.data(), static_cast(audio.size()), + filters.data(), 201, mels, 400, 160, 30, 16000); +} + +std::vector tone(int samples) { + std::vector audio(samples); + constexpr double pi2 = 6.28318530717958647692; + for (int i = 0; i < samples; ++i) { + audio[i] = static_cast(0.25 * std::cos(pi2 * 3.0 * i / 400.0) + + 0.15 * std::sin(pi2 * 7.0 * i / 400.0)); + } + return audio; +} + +// Goldens from Transformers 5.2.0 WhisperFeatureExtractor with the sparse +// filterbank above and n_fft=400, hop_length=160, sampling_rate=16000, +// chunk_length=30. The signal has nonzero endpoints so center padding is visible. +void test_center_padding(int mels) { + const auto short_result = extract(tone(16000), mels); + check_values(short_result, {{0, 0, 1.22718477F}, + {2, 0, 1.72748566F}, + {6, 0, 0.91715926F}, + {0, 1, 0.67869687F}, + {0, 2999, -0.27251434F}}); + + auto audio = tone(480000); + const auto full_result = extract(audio, mels); + check_values(full_result, {{0, 0, 1.22718477F}, + {0, 2999, 0.70422518F}, + {2, 2999, 1.69706297F}, + {6, 2999, 1.58521879F}}); + + audio.resize(480160, 1.0F); + check(extract(audio, mels).data == full_result.data, + "samples beyond the 30-second chunk do not affect its reflection"); +} + +void test_discarded_frame_does_not_set_log_floor(int mels) { + std::vector audio(480000, 0.0F); + for (int i = 0; i < 160; ++i) { + audio[480000 - 160 + i] = static_cast(i + 1) / 320.0F; + } + const auto result = extract(audio, mels); + check_values(result, {{0, 0, -0.40593076F}, + {79, 1000, -0.40593076F}, + {0, 2999, 1.59406924F}, + {2, 2999, 1.10322535F}, + {6, 2999, 0.74954677F}}); +} + +void test_empty_and_single_sample(int mels) { + const auto empty = extract({}, mels); + check(empty.n_mels == mels && empty.n_frames == 3000, "empty input keeps output shape"); + check(std::all_of(empty.data.begin(), empty.data.end(), + [](float value) { return value == -1.5F; }), + "empty input produces the silence floor"); + check_values(extract({0.5F}, mels), + {{0, 0, 0.84948498F}, {0, 1, 0.33946735F}, {0, 2999, -1.15051508F}}); +} + +} // namespace + +int main() { + for (const int mels : {80, 128}) { + test_center_padding(mels); + test_discarded_frame_does_not_set_log_floor(mels); + test_empty_and_single_sample(mels); + } + if (failures == 0) { + std::printf("Whisper mel frontend tests passed\n"); + } + return failures == 0 ? 0 : 1; +}