diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b03d3a..a1ffde4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,16 @@ All notable changes to SKaiNET-audio are documented here. The format follows ## [Unreleased] +### Changed + +- **`audio-mel`: `MelSpectrogram` uses an FFT instead of the direct DFT** — a mixed-radix FFT + (radix 4, 2 and 5 butterflies in closed form, 3 and larger factors generically), so Whisper's + 400-point frame (4·4·5·5) needs no power-of-two size. Each mel filter is summed over its non-zero + band only. A 4 s Whisper window takes 5.8 ms instead of 38 ms on the JVM, and on an arm32 + Android device the log-mel drops from ≈ 0.88 s to ≈ 0.2 s; the 640 KB cos/sin tables are gone. + Output against the `whisper.log_mel_spectrogram` golden: cosine 0.99999999 (was 0.99999998). + `MelSpectrogram` stays safe to share across threads (the FFT's buffers are per call). + ## [0.2.1] - 2026-09-13 ### Changed diff --git a/audio-mel/src/commonMain/kotlin/sk/ainet/audio/mel/MelSpectrogram.kt b/audio-mel/src/commonMain/kotlin/sk/ainet/audio/mel/MelSpectrogram.kt index 2126f4e..46a2e67 100644 --- a/audio-mel/src/commonMain/kotlin/sk/ainet/audio/mel/MelSpectrogram.kt +++ b/audio-mel/src/commonMain/kotlin/sk/ainet/audio/mel/MelSpectrogram.kt @@ -5,7 +5,6 @@ import sk.ainet.audio.core.FeatureExtractor import kotlin.math.PI import kotlin.math.cos import kotlin.math.log10 -import kotlin.math.sin /** * Mel front-end parameters. Presets cover the ASR models this family supports; anything else is a @@ -45,7 +44,7 @@ public enum class MelNorm { } /** - * Log-mel spectrogram, pure Kotlin (reflect-pad ±nFft/2, periodic Hann, nFft-point DFT, mel + * Log-mel spectrogram, pure Kotlin (reflect-pad ±nFft/2, periodic Hann, nFft-point FFT, mel * filterbank, then [MelNorm]). With `MelConfig.whisper()` and Whisper's shipped filterbank this is * verified against `whisper.log_mel_spectrogram` (cosine ≈ 1.0) — the generalized port of the * proven `WhisperMel`. @@ -66,12 +65,16 @@ public class MelSpectrogram( init { require(filters.size == nMels * nBins) { "filterbank ${filters.size} != $nMels*$nBins" } } + // Each mel filter is non-zero on a narrow band of bins only: summing over [bandStart, bandEnd) skips exact zeros, + // so the result is the same as over all bins, at a fraction of the cost. + private val bandStart = IntArray(nMels) { m -> (0 until nBins).firstOrNull { filters[m * nBins + it] != 0f } ?: 0 } + private val bandEnd = IntArray(nMels) { m -> ((nBins - 1 downTo 0).firstOrNull { filters[m * nBins + it] != 0f } ?: -1) + 1 } + // periodic Hann (denominator nFft, matching torch.hann_window(periodic=true)) private val hann = FloatArray(nFft) { (0.5 - 0.5 * cos(2.0 * PI * it / nFft)).toFloat() } - // DFT tables: cos/sin(2π k n / nFft); sign irrelevant since we take magnitude². - private val cosT = Array(nBins) { k -> FloatArray(nFft) { n -> cos(2.0 * PI * k * n / nFft).toFloat() } } - private val sinT = Array(nBins) { k -> FloatArray(nFft) { n -> sin(2.0 * PI * k * n / nFft).toFloat() } } + // mixed-radix FFT: nFft = 400 (Whisper) is 4·4·5·5, so no power-of-two restriction + private val spectrum = PowerSpectrum(nFft) override fun extract(pcm16k: FloatArray): FeatureFrames = logMel(pcm16k) @@ -84,18 +87,14 @@ public class MelSpectrogram( val mel = FloatArray(nMels * frames) val frame = FloatArray(nFft) val power = FloatArray(nBins) + val workspace = spectrum.workspace() for (f in 0 until frames) { val off = f * hop for (n in 0 until nFft) frame[n] = padded[off + n] * hann[n] - for (k in 0 until nBins) { - var re = 0f; var im = 0f - val ck = cosT[k]; val sk = sinT[k] - for (n in 0 until nFft) { re += frame[n] * ck[n]; im += frame[n] * sk[n] } - power[k] = re * re + im * im - } + spectrum.power(frame, power, workspace) for (m in 0 until nMels) { var s = 0f; val base = m * nBins - for (k in 0 until nBins) s += filters[base + k] * power[k] + for (k in bandStart[m] until bandEnd[m]) s += filters[base + k] * power[k] mel[m * frames + f] = s } } diff --git a/audio-mel/src/commonMain/kotlin/sk/ainet/audio/mel/PowerSpectrum.kt b/audio-mel/src/commonMain/kotlin/sk/ainet/audio/mel/PowerSpectrum.kt new file mode 100644 index 0000000..5298460 --- /dev/null +++ b/audio-mel/src/commonMain/kotlin/sk/ainet/audio/mel/PowerSpectrum.kt @@ -0,0 +1,128 @@ +package sk.ainet.audio.mel + +import kotlin.math.PI +import kotlin.math.cos +import kotlin.math.sin + +/** + * Power spectrum `|X_k|²`, `k = 0..n/2`, of a real frame of length [n], by a mixed-radix FFT: `n` is factored into + * 4s, 2s, 3s and 5s (any remaining factor is handled by a direct DFT of that size), so Whisper's 400-point frame + * (4·4·5·5) and NeMo's 512-point frame both run in `O(n·Σfactors)` instead of the direct DFT's `O(n²/2)`. + * + * Recursive decimation in time over a precomputed twiddle table `e^(−2πi·j/n)`; twiddles are computed in double + * precision. The instance is immutable and can be shared; the buffers of one transform live in a [Workspace]. + */ +internal class PowerSpectrum(private val n: Int) { + init { require(n >= 1) { "n must be positive: $n" } } + + private val factors: IntArray = factor(n) + private val wRe = FloatArray(n) { cos(-2.0 * PI * it / n).toFloat() } + private val wIm = FloatArray(n) { sin(-2.0 * PI * it / n).toFloat() } + private val maxRadix = factors.maxOrNull() ?: 1 + + /** The buffers of one transform: one per thread, reusable across frames. */ + class Workspace internal constructor(n: Int, radix: Int) { + internal val re = FloatArray(n) + internal val im = FloatArray(n) + internal val tRe = FloatArray(radix) + internal val tIm = FloatArray(radix) + } + + fun workspace(): Workspace = Workspace(n, maxRadix) + + /** Writes `|X_k|²` of [frame] (length [n]) to [out] (length ≥ `n/2 + 1`). */ + fun power(frame: FloatArray, out: FloatArray, ws: Workspace = workspace()) { + require(frame.size == n) { "frame has ${frame.size} samples, this transform takes $n" } + transform(frame, 0, 1, n, 0, 0, ws) + val re = ws.re; val im = ws.im + for (k in 0..n / 2) out[k] = re[k] * re[k] + im[k] * im[k] + } + + /** + * Transform of the [len] samples `x[inOff + j·stride]` into `re/im[outOff ..< outOff + len]`: the `p` interleaved + * sub-sequences are transformed recursively into consecutive blocks of `m = len/p`, then combined. + */ + private fun transform(x: FloatArray, inOff: Int, stride: Int, len: Int, outOff: Int, level: Int, ws: Workspace) { + val re = ws.re; val im = ws.im; val tRe = ws.tRe; val tIm = ws.tIm + if (len == 1) { re[outOff] = x[inOff]; im[outOff] = 0f; return } + val p = factors[level] + val m = len / p + for (q in 0 until p) transform(x, inOff + q * stride, stride * p, m, outOff + q * m, level + 1, ws) + + val step = n / len // twiddle W_len^j = w[j·step] + val rot = n / p // W_p^j = w[j·rot] + for (k in 0 until m) { + // t_q = W_len^{q·k} · Y_q[k] + var wi = 0 + for (q in 0 until p) { + val yr = re[outOff + q * m + k]; val yi = im[outOff + q * m + k] + tRe[q] = yr * wRe[wi] - yi * wIm[wi] + tIm[q] = yr * wIm[wi] + yi * wRe[wi] + wi += k * step + } + // X[k + r·m] = Σ_q t_q · W_p^{q·r}: closed forms for 2, 4 and 5, the plain sum otherwise + val o = outOff + k + when (p) { + 2 -> { + re[o] = tRe[0] + tRe[1]; im[o] = tIm[0] + tIm[1] + re[o + m] = tRe[0] - tRe[1]; im[o + m] = tIm[0] - tIm[1] + } + 4 -> { + val a0r = tRe[0] + tRe[2]; val a0i = tIm[0] + tIm[2] + val a1r = tRe[0] - tRe[2]; val a1i = tIm[0] - tIm[2] + val a2r = tRe[1] + tRe[3]; val a2i = tIm[1] + tIm[3] + val a3r = tRe[1] - tRe[3]; val a3i = tIm[1] - tIm[3] + re[o] = a0r + a2r; im[o] = a0i + a2i + re[o + m] = a1r + a3i; im[o + m] = a1i - a3r // a1 − i·a3 + re[o + 2 * m] = a0r - a2r; im[o + 2 * m] = a0i - a2i + re[o + 3 * m] = a1r - a3i; im[o + 3 * m] = a1i + a3r // a1 + i·a3 + } + 5 -> { + val s7r = tRe[1] + tRe[4]; val s7i = tIm[1] + tIm[4] + val s10r = tRe[1] - tRe[4]; val s10i = tIm[1] - tIm[4] + val s8r = tRe[2] + tRe[3]; val s8i = tIm[2] + tIm[3] + val s9r = tRe[2] - tRe[3]; val s9i = tIm[2] - tIm[3] + re[o] = tRe[0] + s7r + s8r; im[o] = tIm[0] + s7i + s8i + val s5r = tRe[0] + s7r * C1 + s8r * C2; val s5i = tIm[0] + s7i * C1 + s8i * C2 + val s6r = s10i * S1 + s9i * S2; val s6i = -s10r * S1 - s9r * S2 + re[o + m] = s5r + s6r; im[o + m] = s5i + s6i // s6 = −i·(S1·s10 + S2·s9) + re[o + 4 * m] = s5r - s6r; im[o + 4 * m] = s5i - s6i + val s11r = tRe[0] + s7r * C2 + s8r * C1; val s11i = tIm[0] + s7i * C2 + s8i * C1 + val s12r = -s10i * S2 + s9i * S1; val s12i = s10r * S2 - s9r * S1 + re[o + 2 * m] = s11r - s12r; im[o + 2 * m] = s11i - s12i // s12 = i·(S2·s10 − S1·s9) + re[o + 3 * m] = s11r + s12r; im[o + 3 * m] = s11i + s12i + } + else -> for (r in 0 until p) { + var sr = 0f; var si = 0f + var j = 0 // (q·r mod p)·rot, advanced without a division + for (q in 0 until p) { + sr += tRe[q] * wRe[j] - tIm[q] * wIm[j] + si += tRe[q] * wIm[j] + tIm[q] * wRe[j] + j += r * rot; if (j >= n) j -= n + } + re[o + r * m] = sr + im[o + r * m] = si + } + } + } + } + + internal companion object { + // W_5 = e^(−2πi/5): C1 + i·(−S1), W_5² = C2 + i·(−S2) + private val C1 = cos(2.0 * PI / 5).toFloat() + private val S1 = sin(2.0 * PI / 5).toFloat() + private val C2 = cos(4.0 * PI / 5).toFloat() + private val S2 = sin(4.0 * PI / 5).toFloat() + + /** `n` as a product of 4s, 2s, 3s, 5s and whatever remains (prime factors > 5, each its own stage). */ + fun factor(n: Int): IntArray { + val out = ArrayList() + var rest = n + for (p in intArrayOf(4, 2, 3, 5)) while (rest % p == 0) { out += p; rest /= p } + var d = 7 + while (rest > 1 && d * d <= rest) { while (rest % d == 0) { out += d; rest /= d }; d += 2 } + if (rest > 1) out += rest + return out.toIntArray() + } + } +} diff --git a/audio-mel/src/commonTest/kotlin/sk/ainet/audio/mel/PowerSpectrumTest.kt b/audio-mel/src/commonTest/kotlin/sk/ainet/audio/mel/PowerSpectrumTest.kt new file mode 100644 index 0000000..fbe098d --- /dev/null +++ b/audio-mel/src/commonTest/kotlin/sk/ainet/audio/mel/PowerSpectrumTest.kt @@ -0,0 +1,68 @@ +package sk.ainet.audio.mel + +import kotlin.math.PI +import kotlin.math.abs +import kotlin.math.cos +import kotlin.math.sin +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertTrue + +class PowerSpectrumTest { + + /** `|X_k|²` by the direct DFT, in double precision — the reference. */ + private fun directPower(x: FloatArray): DoubleArray { + val n = x.size + return DoubleArray(n / 2 + 1) { k -> + var re = 0.0; var im = 0.0 + for (j in 0 until n) { val a = -2.0 * PI * k * j / n; re += x[j] * cos(a); im += x[j] * sin(a) } + re * re + im * im + } + } + + private fun assertMatchesDirect(n: Int, seed: Int = n) { + val rnd = Random(seed) + val x = FloatArray(n) { rnd.nextFloat() * 2f - 1f } + val want = directPower(x) + val got = FloatArray(n / 2 + 1).also { PowerSpectrum(n).power(x, it) } + val scale = want.maxOrNull()!!.coerceAtLeast(1e-12) + for (k in want.indices) { + val err = abs(got[k] - want[k]) / scale + assertTrue(err < 1e-5, "n=$n k=$k: fft ${got[k]} vs dft ${want[k]} (relative error $err)") + } + } + + @Test fun whispersFrameSize() = assertMatchesDirect(400) // 4·4·5·5 + @Test fun nemosFrameSize() = assertMatchesDirect(512) // 4·4·4·4·2 + @Test fun everyRadixPath() { for (n in listOf(1, 2, 3, 4, 5, 6, 8, 9, 12, 15, 16, 25, 30, 60, 100, 210, 256)) assertMatchesDirect(n) } + @Test fun primeAndRepeatedLargeFactors() { for (n in listOf(7, 49, 97, 397)) assertMatchesDirect(n) } + + @Test + fun aPureToneLandsInItsBin() { + val n = 400 + val x = FloatArray(n) { cos(2.0 * PI * 25 * it / n).toFloat() } + val p = FloatArray(n / 2 + 1).also { PowerSpectrum(n).power(x, it) } + val peak = p.indices.maxByOrNull { p[it] }!! + assertTrue(peak == 25 && abs(p[25] - (n / 2.0) * (n / 2.0)) / ((n / 2.0) * (n / 2.0)) < 1e-5, "peak at $peak = ${p[peak]}") + } + + @Test + fun oneInstanceServesSeveralWorkspaces() { + val fft = PowerSpectrum(400) + val x = FloatArray(400) { sin(0.1 * it).toFloat() } + val a = FloatArray(201).also { fft.power(x, it, fft.workspace()) } + val b = FloatArray(201).also { fft.power(x, it, fft.workspace()) } + assertContentEquals(a, b) + } + + @Test + fun factorsPreferRadixFour() { + assertContentEquals(intArrayOf(4, 4, 5, 5), PowerSpectrum.factor(400)) + assertContentEquals(intArrayOf(4, 4, 4, 4, 2), PowerSpectrum.factor(512)) + assertContentEquals(intArrayOf(2, 3, 5, 7), PowerSpectrum.factor(210)) + assertContentEquals(intArrayOf(7, 7), PowerSpectrum.factor(49)) + assertContentEquals(intArrayOf(397), PowerSpectrum.factor(397)) + assertContentEquals(intArrayOf(), PowerSpectrum.factor(1)) + } +}