musetric/beat-this-onnx

Beat This! — beat & downbeat tracking (ONNX, WebGPU)

ONNX export of Beat This!, the JKU beat/downbeat tracker (checkpoint final0), packaged for the musetric packages/ai runtime on onnxruntime-web's WebGPU execution provider.

The graph covers the network only:

beat_this.onnx  spect [windows, frames, 128] → beat, downbeat [windows, frames]

Feature extraction is the host's job, and mel-filterbank.bin is what it needs to do it: the host computes the log-mel spectrogram on the GPU (STFT + this filterbank + log1p) and feeds windows straight into the graph as GPU buffers, never round-tripping features through the CPU.

Feed one window per call. Batching every window into a single call materializes an attention tensor of windows × 32 × 1500 × 1500 floats — about 2.9 GB on a five-minute track. This is not a tuning knob; it is why the graph takes a window axis at all.

The export substitutes no algorithm. The one construct that cannot be traced — PartialFTTransformer.forward's b = len(x), which would freeze the window count into the graph — is re-expressed with x.shape, and is bit-identical to the original in torch. So the graph carries no approximation and its parity does not depend on the material.

Intended uses & limitations

Intended:

• Beat and downbeat times for a music track, as a stage in an audio pipeline; tempo and meter are derived from them host-side.
• GPU inference through onnxruntime-web. The wasm EP also works and gives identical results, at roughly 12× the wall clock.

Out of scope:

• Use in other training frameworks — this is an inference-only export.
• The madmom DBN postprocessor. This export targets the reference's minimal postprocessing, which is what Beat This! is designed around.

Limitations:

• The host owns feature extraction, and the contract is exact. 22050 Hz mono; channels downmixed as the arithmetic mean, matching the reference Audio2Frames; STFT with n_fft 1024, hop 441, a periodic Hann window, center=True with reflect padding; magnitude divided by sqrt(n_fft) (torchaudio's normalized='frame_length'); projected onto mel-filterbank.bin; then log1p(1000 · x). config.json records all of it.
• ffmpeg -ac 1 is the wrong downmix. It is an energy-preserving rematrix and comes out √2 louder; because features are log1p(1000 · mel), that gain is not a constant offset and will move beats. config.json says "downmix": "mean".
• The host also owns the chunk layout. The graph does not know about chunkSize, borderSize or the shifted final window. Reproduce split_piece / aggregate_prediction (keep_first) or results drift at window seams.
• Validate on the material fed in production. For this pipeline that is the instrumental stem.
• Training-data provenance of the upstream checkpoint is not documented here.

How to use

import * as ort from 'onnxruntime-web/webgpu';

const session = await ort.InferenceSession.create('beat_this.onnx', {
  executionProviders: ['webgpu'],
});

// window: Float32Array(1500 * 128), one log-mel chunk (see config.json).
const { beat, downbeat } = await session.run({
  spect: new ort.Tensor('float32', window, [1, 1500, 128]),
});

mel-filterbank.bin is a raw row-major float32 matrix, [513, 128] — load it with new Float32Array(await (await fetch(url)).arrayBuffer()) and project magnitudes onto it.

Logits are peak-picked host-side: keep maxima within ±3 frames (peakKernel 7) whose logit exceeds peakThreshold 0, merge peaks ≤ deduplicateWidth frames apart, convert frames to seconds at fps 50, then snap each downbeat to its nearest beat. See the musetric packages/ai host code (runtime/rhythm/, rhythm/beatPeaks.ts) for the full pipeline.

Files

File	Size	SHA256
beat_this.onnx	83,143,431 B	078572af6ca47741e06a82d09525d13c793eaa8e311a8cf15e831dcd7e73f218
mel-filterbank.bin	262,656 B	1ee975d96f44ccf2c3bfe37825c1c1f0b089f5703c7a12a84b1f0a3bce004533
config.json	1,024 B	56cc961ddc588c57787c20c01ec6ab483b23af1049e65bd33d599a81803acd69

mel-filterbank.bin is torchaudio's own MelScale.fb written verbatim, so a host reuses the reference filterbank instead of reimplementing the slaney mel scale. The analysis window is deliberately not shipped: it is a plain periodic Hann, 0.5 - 0.5·cos(2πn/N), and the export refuses to run if the reference window ever stops matching that.

Signature — float32 weights, opset ai.onnx 17:

Tensor	Type	Shape	Meaning
spect (in)	float32	[windows, frames, 128]	one 1500-frame log-mel window per call
beat (out)	float32	[windows, frames]	beat logits; > 0 is a candidate
downbeat (out)	float32	[windows, frames]	downbeat logits

Validation

The graph, against the PyTorch File2Beats reference sharing its log-mel and minimal postprocessing, over 20 instrumental stems (the production input, 10 s–285 s):

Metric	Value
tracks with identical beat and downbeat counts	20/20
max beat / downbeat time difference	0.0 s

Beat times are bit-identical: logits agree to ~1e-5 and peak picking quantizes to 20 ms frames, so the residual cannot move a beat.

The WebGPU front end, against torchaudio's LogMelSpect on the same waveforms:

Metric	Value
frame counts	20/20 exact
max absolute difference	3.2e-04 (relative to peak: 3.8e-05)
mean absolute difference	~1e-06

That residual is float32 GPU arithmetic, not a different transform, and it changes nothing downstream: the WebGPU path reproduces the wasm path's results byte for byte on 20/20 tracks.

End to end, the shipped host (ffmpeg mean-downmix decode + WebGPU) against the Python CLI it replaces (torchaudio + soxr), over the same 20 stems:

Metric	Value
bpm — the user-visible tempo	20/20 identical
meter	20/20 identical
beat F-measure (±70 ms)	mean 0.9982, median 1.0000, min 0.9898
downbeat F-measure (±70 ms)	mean 0.9962, median 1.0000, min 0.9722
beats differing over the whole set	20 missed + 9 extra of 7,853 (0.37 %)

Neither the graph nor the front end causes that, and neither does gain: the mean downmix matches the reference to a ratio of 1.0000 and the two decodes correlate at 0.99998 with zero sample shift. What is left is the resampler — ffmpeg's swr versus soxr — and peak picking turns that into a hard decision, so a handful of marginal beats near the logit > 0 threshold flip. Tempo and meter are medians over hundreds of beats and absorb it. Feeding the reference's own decode through this stack reproduces the CLI exactly, which is what pins the residual on the decoder.

Re-run the parity gate on the exact published bytes before relying on it.

Source & lineage

Code license and weight license are separate; ONNX conversion does not change the weight license. Documented only as far as it is verifiable.

• Architecture: log-mel front end (n_fft 1024, hop 441, 128 mels, log1p(1000 · x)) → a conv stem and three frontend blocks, each a partial transformer attending across frequency then time → six rotary transformer layers (dim 512, 16 heads) → a sum head emitting beat and downbeat logits. 20.25 M parameters.
• Reference implementation and weights: CPJKU/beat_this, MIT (per its LICENSE, Copyright (c) 2024 Institute of Computational Perception, JKU Linz, Austria), consumed as the beat-this package (1.1.0).
• Paper: Foscarin, Schlüter, Widmer — Beat this! Accurate beat tracking without DBN postprocessing, ISMIR 2024.
• Checkpoint: final0 — beat_this-final0.ckpt, 81,058,141 B, sha256 8c328b45f59d8dd3dff219253ff6a8d6482be57d0133a29140e2febbf8eb8331 — fetched at export time from the upstream checkpoint host via torch.hub, as the package's load_checkpoint does. Upstream serves it from a cloud share rather than a content-addressed URL, so the sha256 above is the pin.
• Export tooling: scripts/onnx/beat_this in musetric-toolkit; see its thirdPartyNotices.md.
• Host runtime: packages/ai in musetric.

This export preserves the upstream MIT license; we do not claim authorship of the original weights.

License

MIT, inherited from the upstream weights.