add onnx and metadata
This commit is contained in:
Overside srl committed
1 parent
98d9fa24cf
commit
5f3f915258
18 files changed
+2229
-78
No files matched your search
@@ -0,0 +1,330 @@
|
||||
"""
|
||||
Designed (non-trained) genre CNN for the ldcc bridge.
|
||||
|
||||
The weights are WRITTEN, not learned: there is no dataset and no gradient descent.
|
||||
Everything is deterministic and reproducible.
|
||||
|
||||
Design
|
||||
------
|
||||
1. Conv layer with fixed, interpretable filters: eight uniform averages over
|
||||
contiguous 16-band mel groups. No randomness, no training.
|
||||
2. Temporal statistics per band group: mean, standard deviation and mean absolute
|
||||
frame-to-frame difference (a cheap onset/flux proxy). 8 x 3 = 24 features.
|
||||
3. Linear head obtained in CLOSED FORM (ridge) so that the 24 features of five
|
||||
reference sounds map to five one-hot vectors. The references are synthesized
|
||||
right here, not recorded, so no external data and no licensing issues.
|
||||
|
||||
What this is and is not
|
||||
-----------------------
|
||||
- It is: a complete, valid ONNX model that satisfies the bridge contract, with
|
||||
weights derived from an explicit procedure and inspectable in the source.
|
||||
- It is not: a trained classifier. Accuracy on real music will be roughly the
|
||||
accuracy of the five reference prototypes. Train with tools/train_genre_model.py
|
||||
on real labelled audio to beat it.
|
||||
|
||||
Usage
|
||||
-----
|
||||
python tools/design_cnn.py --output-dir bin
|
||||
python tools/design_cnn.py --validate # checks unseen renders
|
||||
|
||||
Contract (identical to the bridge):
|
||||
input [1, 1, 128, 300] float32 = (batch, channel, mel bands, frames)
|
||||
output [1, 5] float32 = probabilities, label order LABELS
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||
import train_genre_model as tgm # noqa: E402 (front-end + contract, single source)
|
||||
|
||||
SR = 48000
|
||||
CROP_SECONDS = 3.0
|
||||
# Normalization constants are a design choice: they must match the json shipped
|
||||
# with the model, because the bridge applies them before the tensor.
|
||||
MEAN = -4.0
|
||||
STD = 4.0
|
||||
N_GROUPS = tgm.N_MELS // 16 # 128 / 16 = 8 mel groups
|
||||
N_FEATURES = N_GROUPS * 3 # mean, std, flux
|
||||
BPM_RANGE = {"latin": (85, 105), "electronic": (115, 135), "hiphop": (80, 95),
|
||||
"rock": (125, 150), "classical": (70, 80)}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Synthesis of reference sounds (deterministic, no external data)
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _add(buf, i0, wave):
|
||||
"""Somma un burst nel buffer, troncando oltre la fine (niente broadcast negativo)."""
|
||||
if i0 >= buf.size:
|
||||
return
|
||||
m = min(wave.size, buf.size - i0)
|
||||
buf[i0:i0 + m] += wave[:m]
|
||||
|
||||
|
||||
def _kick(buf, sr, t0, amp, rng):
|
||||
n = int(0.35 * sr)
|
||||
i0 = int(t0 * sr)
|
||||
t = np.arange(n) / sr
|
||||
wave = np.sin(2 * np.pi * (55 * t + (1 - np.exp(-t * 40)) / 40)) * np.exp(-t * 16) * amp
|
||||
_add(buf, i0, wave)
|
||||
|
||||
|
||||
|
||||
def _snare(buf, sr, t0, amp, rng):
|
||||
n = int(0.15 * sr)
|
||||
i0 = int(t0 * sr)
|
||||
t = np.arange(n) / sr
|
||||
noise = (rng.standard_normal(n) * 2 - 1) * np.exp(-t * 28)
|
||||
body = np.sin(2 * np.pi * 190 * t) * np.exp(-t * 30) * 0.4
|
||||
wave = (noise + body) * amp
|
||||
_add(buf, i0, wave)
|
||||
|
||||
|
||||
|
||||
def _hat(buf, sr, t0, amp, rng):
|
||||
n = int(0.04 * sr)
|
||||
i0 = int(t0 * sr)
|
||||
t = np.arange(n) / sr
|
||||
noise = rng.standard_normal(n)
|
||||
hp = np.diff(noise, prepend=0.0)
|
||||
wave = hp * np.exp(-t * 90) * amp
|
||||
_add(buf, i0, wave)
|
||||
|
||||
|
||||
|
||||
def _chord(buf, sr, t0, dur, amp):
|
||||
n = int(dur * sr)
|
||||
i0 = int(t0 * sr)
|
||||
t = np.arange(n) / sr
|
||||
env = np.exp(-t * 3.0)
|
||||
v = np.zeros(n)
|
||||
for root in (110.0, 146.83, 164.81):
|
||||
for h in range(1, 21):
|
||||
v += np.sin(2 * np.pi * root * h * t + h) / h
|
||||
wave = np.tanh(v * 6.0) * env * amp / 10.0
|
||||
_add(buf, i0, wave)
|
||||
|
||||
|
||||
|
||||
def _tones(buf, sr, seconds, rng):
|
||||
notes = (220.0, 261.63, 329.63, 392.0, 329.63, 261.63)
|
||||
t = 0.0
|
||||
i = 0
|
||||
while t < seconds:
|
||||
f = notes[i % len(notes)]
|
||||
dur = 1.6 + rng.random() * 1.2
|
||||
amp = 0.12 + rng.random() * 0.5
|
||||
n = int(dur * sr)
|
||||
i0 = int(t * sr)
|
||||
tt = np.arange(n) / sr
|
||||
env = np.exp(-tt * 1.2) * (1 - np.exp(-tt * 60))
|
||||
wave = (np.sin(2 * np.pi * f * tt) + 0.3 * np.sin(4 * np.pi * f * tt)) * env * amp
|
||||
_add(buf, i0, wave)
|
||||
t += dur * 0.85
|
||||
i += 1
|
||||
|
||||
|
||||
def synth(kind: str, bpm: float, seconds: float, seed: int) -> np.ndarray:
|
||||
rng = np.random.default_rng(seed)
|
||||
sr = SR
|
||||
buf = np.zeros(int(sr * seconds), dtype=np.float64)
|
||||
sixteenth = 60.0 / bpm / 4.0
|
||||
kick, snare, hat, chord = [], [], [], False
|
||||
if kind == "latin":
|
||||
kick, snare = [0, 8], [0, 3, 6, 8, 11, 14]
|
||||
elif kind == "electronic":
|
||||
kick, hat = [0, 4, 8, 12], [2, 6, 10, 14]
|
||||
elif kind == "hiphop":
|
||||
kick, snare, hat = [0, 6], [4, 12], [2, 6, 10, 14]
|
||||
elif kind == "rock":
|
||||
kick, snare, hat = [0, 6, 8], [4, 12], [0, 2, 4, 6, 8, 10, 12, 14]
|
||||
chord = True
|
||||
elif kind == "classical":
|
||||
_tones(buf, sr, seconds, rng)
|
||||
return np.clip(buf, -1, 1).astype(np.float32)
|
||||
else:
|
||||
raise ValueError(kind)
|
||||
|
||||
bar = 0
|
||||
while bar * 16 * sixteenth < seconds:
|
||||
base = bar * 16 * sixteenth
|
||||
for s in kick:
|
||||
_kick(buf, sr, base + s * sixteenth, 0.85, rng)
|
||||
for s in snare:
|
||||
_snare(buf, sr, base + s * sixteenth, 0.45, rng)
|
||||
for s in hat:
|
||||
_hat(buf, sr, base + s * sixteenth, 0.18, rng)
|
||||
if chord:
|
||||
_chord(buf, sr, base, sixteenth * 8, 0.30)
|
||||
bar += 1
|
||||
return np.clip(buf, -1, 1).astype(np.float32)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Fixed convolution + temporal statistics
|
||||
# --------------------------------------------------------------------------- #
|
||||
def band_group_weights() -> torch.Tensor:
|
||||
"""Uno filtro per gruppo di bande: kernel [8,1,128,1], canale g attivo (1/16) sulle
|
||||
righe mel 16g..16g+15. Deterministico e leggibile, nessuna inizializzazione casuale."""
|
||||
w = torch.zeros(N_GROUPS, 1, tgm.N_MELS, 1, dtype=torch.float32)
|
||||
for g in range(N_GROUPS):
|
||||
w[g, 0, g * 16:(g + 1) * 16, 0] = 1.0 / 16.0
|
||||
return w
|
||||
|
||||
|
||||
class DesignedGenreCNN(nn.Module):
|
||||
"""Conv(fixed) -> (mean, std, flux) -> Linear(fitted)."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.bands = nn.Conv2d(1, N_GROUPS, (tgm.N_MELS, 1), stride=(1, 1), bias=False)
|
||||
self.bands.weight = nn.Parameter(band_group_weights(), requires_grad=False)
|
||||
self.head = nn.Linear(N_FEATURES, len(tgm.LABELS))
|
||||
|
||||
def pool(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# x: [N, G, 1, T]
|
||||
x = x.squeeze(2) # [N, G, T]
|
||||
mean = x.mean(dim=2) # [N, G]
|
||||
var = torch.clamp((x * x).mean(dim=2) - mean * mean, min=0.0)
|
||||
std = torch.sqrt(var + 1e-8) # [N, G]
|
||||
flux = torch.abs(x[:, :, 1:] - x[:, :, :-1]).mean(dim=2)
|
||||
return torch.cat([mean, std, flux], dim=1) # [N, 24]
|
||||
|
||||
def embed(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# Invarianza al livello: il bridge normalizza con una costante globale, ma il
|
||||
# volume di riproduzione sposta tutto lo spettro. Sottrarre la media della
|
||||
# finestra rende il modello insensibile a un guadagno costante.
|
||||
return self.pool(self.bands(x - x.mean()))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.softmax(self.head(self.embed(x)), dim=1)
|
||||
|
||||
|
||||
def features_of(audio: np.ndarray) -> np.ndarray:
|
||||
"""Run the bridge front-end on a raw clip and return the 24 pooled features."""
|
||||
model = DesignedGenreCNN()
|
||||
ten = torch.from_numpy(tgm.tensor_frontend(audio, SR, MEAN, STD))
|
||||
with torch.no_grad():
|
||||
return model.embed(ten).numpy()[0]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Closed-form head
|
||||
# --------------------------------------------------------------------------- #
|
||||
def fit_head(features: np.ndarray, gain: float = 2.5) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Minimum-norm least-squares fit of the 5 reference rows to one-hot targets.
|
||||
|
||||
Closed form, deterministic, no iterative training. `gain` sharpens the logits:
|
||||
the minimum-norm solution is deliberately small, and without a gain the softmax
|
||||
would stay close to uniform.
|
||||
"""
|
||||
x = np.concatenate([features, np.ones((features.shape[0], 1), dtype=np.float64)], axis=1)
|
||||
y = np.eye(len(tgm.LABELS))
|
||||
solution, *_ = np.linalg.lstsq(x, y, rcond=None) # [features+1, classes]
|
||||
w = solution[:-1].T.astype(np.float32) * gain
|
||||
b = solution[-1].astype(np.float32) * gain
|
||||
return w, b
|
||||
|
||||
|
||||
def build_references(seed: int = 11) -> tuple[np.ndarray, list[str]]:
|
||||
rows = []
|
||||
used = []
|
||||
for label in tgm.LABELS:
|
||||
lo, hi = BPM_RANGE[label]
|
||||
bpm = 0.5 * (lo + hi)
|
||||
audio = synth(label, bpm, 6.0, seed)
|
||||
rows.append(features_of(audio))
|
||||
used.append(f"{label}@{bpm:.0f}bpm")
|
||||
return np.vstack(rows), used
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--output-dir", type=Path, default=Path(__file__).resolve().parents[1] / "bin")
|
||||
parser.add_argument("--validate", action="store_true", help="classify unseen renders")
|
||||
parser.add_argument("--seed", type=int, default=11)
|
||||
args = parser.parse_args()
|
||||
|
||||
features, used = build_references(args.seed)
|
||||
print("PROTOTYPES: " + ", ".join(used), flush=True)
|
||||
w, b = fit_head(features)
|
||||
|
||||
model = DesignedGenreCNN()
|
||||
with torch.no_grad():
|
||||
model.head.weight.copy_(torch.from_numpy(w))
|
||||
model.head.bias.copy_(torch.from_numpy(b))
|
||||
model.eval()
|
||||
|
||||
# Sanity: the references must come back as themselves (head applied to the
|
||||
# pooled features, since `features` is already past the convolution).
|
||||
with torch.no_grad():
|
||||
probs = model.head(torch.from_numpy(features)).softmax(dim=1)
|
||||
print("REFERENCE FIT:")
|
||||
for i, label in enumerate(tgm.LABELS):
|
||||
pred = tgm.LABELS[int(probs[i].argmax())]
|
||||
print(f" {label:11s} -> {pred:11s} p={probs[i].max():.3f}")
|
||||
|
||||
if args.validate:
|
||||
print("UNSEEN RENDERS:")
|
||||
ok = 0
|
||||
total = 0
|
||||
for label in tgm.LABELS:
|
||||
lo, hi = BPM_RANGE[label]
|
||||
for frac in (0.25, 0.75):
|
||||
bpm = lo + (hi - lo) * frac
|
||||
audio = synth(label, bpm, 6.0, args.seed + 100 + total)
|
||||
f = features_of(audio).reshape(1, -1)
|
||||
with torch.no_grad():
|
||||
p = model.head(torch.from_numpy(f)).softmax(dim=1).numpy()[0]
|
||||
pred = tgm.LABELS[int(np.argmax(p))]
|
||||
ok += pred == label
|
||||
total += 1
|
||||
print(f" {label:11s} {bpm:6.1f}bpm seed={args.seed + 100 + total - 1} -> {pred:11s} p={p.max():.3f}")
|
||||
print(f"VALIDATION: {ok}/{total}")
|
||||
if ok < total * 0.6:
|
||||
print("WARNING: too few unseen renders classified correctly; this model is a scaffold.")
|
||||
|
||||
out = args.output_dir
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
onnx_path = out / "genre_model.onnx"
|
||||
json_path = out / "genre_model.json"
|
||||
example = torch.zeros(1, 1, tgm.N_MELS, tgm.FRAMES, dtype=torch.float32)
|
||||
torch.onnx.export(
|
||||
model, example, str(onnx_path),
|
||||
input_names=["logmel"], output_names=["probabilities"],
|
||||
opset_version=17, dynamo=False,
|
||||
)
|
||||
with torch.no_grad():
|
||||
check = model(example).numpy()[0]
|
||||
meta = {
|
||||
"labels": tgm.LABELS,
|
||||
"mean": MEAN,
|
||||
"std": STD,
|
||||
"version": tgm.VERSION,
|
||||
"source": "designed",
|
||||
"frontend": {
|
||||
"fft": 2048, "hann": "periodic", "hop": "sr//100",
|
||||
"padding": "causal-left-N-hop", "power": "abs(rfft)^2",
|
||||
"mel": "HTK", "bands": tgm.N_MELS, "fmin": 30, "fmax": 11000,
|
||||
"area": "divide by full support width in fractional bins", "log_floor": 1e-10,
|
||||
},
|
||||
}
|
||||
json_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
|
||||
print(f"MODEL={onnx_path.resolve()} ({onnx_path.stat().st_size} bytes)")
|
||||
print(f"JSON={json_path.resolve()}")
|
||||
print(f"SANITY: sum={check.sum():.6f} max={check.max():.3f} label={tgm.LABELS[int(check.argmax())]}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,361 @@
|
||||
"""Train/export the fixed ldcc ONNX contract; no-argument invocation runs self-test.
|
||||
|
||||
Python requires only torch, numpy, scipy and onnx. No onnxruntime is needed.
|
||||
Parity: --verify-dump tensor.bin --data SELFTEST_DIR --rate 48000 (or 44100).
|
||||
The fixture is SELFTEST_DIR/fixture.wav; --rate reinterprets samples, NOT resamples.
|
||||
Train and export real data with --data DATA_DIR --output-dir bin.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
import tempfile
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from scipy.io import wavfile
|
||||
from scipy.signal import lfilter
|
||||
import torch
|
||||
from torch import nn
|
||||
import onnx
|
||||
from onnx.reference import ReferenceEvaluator
|
||||
|
||||
LABELS = ["latin", "electronic", "hiphop", "rock", "classical"]
|
||||
N_FFT, N_MELS, FRAMES = 2048, 128, 300
|
||||
FMIN, FMAX, LOG_FLOOR = 30.0, 11000.0, 1e-10
|
||||
VERSION = 1
|
||||
|
||||
|
||||
def read_wav(path: Path) -> tuple[np.ndarray, int]:
|
||||
sr, audio = wavfile.read(path)
|
||||
if sr not in (44100, 48000):
|
||||
raise ValueError(f"{path}: only 44100/48000 Hz supported, got {sr}")
|
||||
if audio.dtype == np.uint8:
|
||||
audio = (audio.astype(np.float64) - 128.0) / 128.0
|
||||
elif np.issubdtype(audio.dtype, np.signedinteger):
|
||||
# scipy returns 24-bit WAV left-justified in int32, matching this scaling.
|
||||
audio = audio.astype(np.float64) / float(2 ** (audio.dtype.itemsize * 8 - 1))
|
||||
elif np.issubdtype(audio.dtype, np.floating):
|
||||
audio = audio.astype(np.float64)
|
||||
else:
|
||||
raise ValueError(f"Unsupported WAV dtype: {audio.dtype}")
|
||||
if audio.ndim == 2:
|
||||
audio = audio.mean(axis=1)
|
||||
if audio.ndim != 1 or not np.all(np.isfinite(audio)):
|
||||
raise ValueError(f"{path}: invalid/nonfinite PCM")
|
||||
# C# ingestion is float32 mono. Promote only AFTER this rounding for FFT parity.
|
||||
return audio.astype(np.float32), sr
|
||||
|
||||
|
||||
@lru_cache(maxsize=2)
|
||||
def mel_bank(sr: int) -> np.ndarray:
|
||||
if sr not in (44100, 48000):
|
||||
raise ValueError("Expected 44100 or 48000 Hz")
|
||||
low = 2595.0 * math.log10(1 + FMIN / 700)
|
||||
high = 2595.0 * math.log10(1 + FMAX / 700)
|
||||
# Fractional edges: do NOT round/floor the FFT-bin locations.
|
||||
edges = np.array([
|
||||
700 * (10 ** ((low + (high - low) * i / (N_MELS + 1)) / 2595) - 1) * N_FFT / sr
|
||||
for i in range(N_MELS + 2)
|
||||
], dtype=np.float64)
|
||||
bins = np.arange(N_FFT // 2 + 1, dtype=np.float64)
|
||||
filters = np.zeros((N_MELS, len(bins)), dtype=np.float64)
|
||||
for m in range(N_MELS):
|
||||
left, center, right = edges[m:m + 3]
|
||||
filters[m] = np.maximum(0, np.minimum((bins - left) / (center - left), (right - bins) / (right - center))) / (right - left)
|
||||
return filters
|
||||
|
||||
|
||||
def stft_manual(x: np.ndarray, sr: int) -> np.ndarray:
|
||||
"""Causal periodic-Hann complex rFFT, shape [1025, floor(len(x)/hop)]."""
|
||||
hop = sr // 100
|
||||
x = np.asarray(x, dtype=np.float32).astype(np.float64)
|
||||
if x.ndim != 1 or not np.isfinite(x).all():
|
||||
raise ValueError("Expected finite mono samples")
|
||||
if len(x) < hop:
|
||||
return np.empty((N_FFT // 2 + 1, 0), dtype=np.complex128)
|
||||
padded = np.pad(x, (N_FFT - hop, 0))
|
||||
frames = np.lib.stride_tricks.sliding_window_view(padded, N_FFT)[::hop]
|
||||
hann = 0.5 - 0.5 * np.cos(2 * np.pi * np.arange(N_FFT, dtype=np.float64) / N_FFT)
|
||||
return np.fft.rfft(frames * hann, axis=1).T
|
||||
|
||||
|
||||
def logmel(x: np.ndarray, sr: int) -> np.ndarray:
|
||||
spectrum = stft_manual(x, sr)
|
||||
# Raw power: no /N and no doubling of positive-frequency bins.
|
||||
power = spectrum.real ** 2 + spectrum.imag ** 2
|
||||
energies = mel_bank(sr) @ power
|
||||
return np.log10(np.maximum(energies, LOG_FLOOR))
|
||||
|
||||
|
||||
def tensor_frontend(x: np.ndarray, sr: int, mean: float, std: float) -> np.ndarray:
|
||||
needed = FRAMES * (sr // 100)
|
||||
if len(x) < needed:
|
||||
raise ValueError(f"Need {needed} samples for 300 frames, got {len(x)}")
|
||||
if not np.isfinite(mean) or not np.isfinite(std) or std <= 0:
|
||||
raise ValueError("Invalid normalization statistics")
|
||||
result = ((logmel(x[:needed], sr) - mean) / std).astype(np.float32)
|
||||
if result.shape != (N_MELS, FRAMES) or not np.isfinite(result).all():
|
||||
raise ValueError("Invalid frontend output")
|
||||
return result[None, None]
|
||||
|
||||
|
||||
def discover(data: Path) -> list[tuple[Path, int]]:
|
||||
found = []
|
||||
for label_id, label in enumerate(LABELS):
|
||||
folder = data / label
|
||||
if folder.exists():
|
||||
found.extend((p.resolve(), label_id) for p in sorted(folder.iterdir()) if p.suffix.lower() == ".wav")
|
||||
if not found:
|
||||
raise ValueError("No WAV files in class subfolders")
|
||||
return found
|
||||
|
||||
|
||||
def split_files(items, seed: int):
|
||||
rng = random.Random(seed)
|
||||
train, val = [], []
|
||||
for label in range(len(LABELS)):
|
||||
group = [item for item in items if item[1] == label]
|
||||
if not group:
|
||||
continue
|
||||
if len(group) < 2:
|
||||
raise ValueError(f"Class {LABELS[label]} needs at least two distinct files")
|
||||
rng.shuffle(group)
|
||||
nval = max(1, round(len(group) * 0.25))
|
||||
val.extend(group[:nval]); train.extend(group[nval:])
|
||||
if set(p for p, _ in train) & set(p for p, _ in val):
|
||||
raise AssertionError("File leakage between train and validation")
|
||||
return train, val
|
||||
|
||||
|
||||
def random_crop(audio, sr, rng):
|
||||
count = FRAMES * (sr // 100)
|
||||
if len(audio) < count:
|
||||
raise ValueError("Training WAV shorter than 3 seconds")
|
||||
offset = int(rng.integers(0, len(audio) - count + 1))
|
||||
return audio[offset:offset + count].copy()
|
||||
|
||||
|
||||
def training_statistics(items, seed):
|
||||
"""Global scalar mean/std on TRAIN FILES ONLY; two reproducible crops/file."""
|
||||
rng = np.random.default_rng(seed)
|
||||
total = squares = 0.0
|
||||
count = 0
|
||||
for path, _ in items:
|
||||
audio, sr = read_wav(path)
|
||||
for _ in range(2):
|
||||
values = logmel(random_crop(audio, sr, rng), sr)
|
||||
total += float(values.sum()); squares += float(np.square(values).sum()); count += values.size
|
||||
mean = total / count
|
||||
std = math.sqrt(max(0, squares / count - mean * mean))
|
||||
if std < 1e-8:
|
||||
raise ValueError("Training set has zero variance")
|
||||
return mean, std
|
||||
|
||||
|
||||
class WavDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, items, mean, std, seed, augment):
|
||||
self.items, self.mean, self.std, self.seed, self.augment = items, mean, std, seed, augment
|
||||
self.rng = np.random.default_rng(seed)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.items)
|
||||
|
||||
def __getitem__(self, index):
|
||||
path, label = self.items[index]
|
||||
audio, sr = read_wav(path)
|
||||
# Validation crop fixed across epochs; train crop randomized on each read.
|
||||
rng = self.rng if self.augment else np.random.default_rng(self.seed + index)
|
||||
x = random_crop(audio, sr, rng)
|
||||
if self.augment:
|
||||
x *= np.float32(10 ** (rng.uniform(-6, 6) / 20))
|
||||
if rng.random() < 0.5:
|
||||
x *= -1
|
||||
return torch.from_numpy(tensor_frontend(x, sr, self.mean, self.std)[0]), label
|
||||
|
||||
|
||||
class GenreCNN(nn.Module):
|
||||
def __init__(self, active_labels):
|
||||
super().__init__()
|
||||
self.features = nn.Sequential(
|
||||
nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
|
||||
nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
|
||||
nn.Conv2d(16, 24, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
|
||||
nn.AdaptiveAvgPool2d((1, 1)),
|
||||
)
|
||||
self.classifier = nn.Linear(24, 5)
|
||||
self.register_buffer("active", torch.tensor([i in active_labels for i in range(5)], dtype=torch.bool))
|
||||
|
||||
def forward(self, x):
|
||||
logits = self.classifier(self.features(x).flatten(1))
|
||||
# Self-test keeps five outputs. Absent classes have EXACT zero probability.
|
||||
return torch.softmax(logits.masked_fill(~self.active, float("-inf")), dim=1)
|
||||
|
||||
|
||||
def epoch(model, loader, optimizer=None):
|
||||
training = optimizer is not None
|
||||
model.train(training)
|
||||
loss_total, correct, count = 0.0, 0, 0
|
||||
with torch.set_grad_enabled(training):
|
||||
for x, labels in loader:
|
||||
if training:
|
||||
optimizer.zero_grad()
|
||||
probs = model(x)
|
||||
loss = -torch.log(probs[torch.arange(len(labels)), labels].clamp_min(1e-9)).mean()
|
||||
if training:
|
||||
loss.backward(); optimizer.step()
|
||||
loss_total += loss.item() * len(labels)
|
||||
correct += (probs.argmax(1) == labels).sum().item(); count += len(labels)
|
||||
return loss_total / count, correct / count
|
||||
|
||||
|
||||
def validate_export(model_path: Path, model: GenreCNN, sample: np.ndarray):
|
||||
exported = onnx.load(str(model_path))
|
||||
onnx.checker.check_model(exported)
|
||||
shape = [dim.dim_value for dim in exported.graph.input[0].type.tensor_type.shape.dim]
|
||||
if shape != [1, 1, 128, 300]:
|
||||
raise AssertionError(f"Bad ONNX shape {shape}")
|
||||
output = ReferenceEvaluator(exported).run(None, {exported.graph.input[0].name: sample})[0]
|
||||
if output.shape != (1, 5) or not np.isfinite(output).all() or np.any(output < 0) or abs(float(output.sum()) - 1) > 1e-5:
|
||||
raise AssertionError(f"Invalid ONNX probabilities: {output}")
|
||||
with torch.no_grad():
|
||||
expected = model(torch.from_numpy(sample)).numpy()
|
||||
difference = float(np.max(np.abs(output - expected)))
|
||||
if difference > 1e-4:
|
||||
raise AssertionError(f"Torch/ONNX difference {difference}")
|
||||
inactive = ~model.active.cpu().numpy()
|
||||
if np.any(output[0, inactive] != 0):
|
||||
raise AssertionError("Unused classes must be exactly zero")
|
||||
print(f"ONNX OK: shape={shape}; sum={output.sum():.9f}; torch_max_abs={difference:.9g}; probs={output.tolist()}", flush=True)
|
||||
|
||||
|
||||
def train(data: Path, output: Path, epochs: int, seed: int, batch_size: int, patience: int):
|
||||
random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
|
||||
items = discover(data)
|
||||
train_items, val_items = split_files(items, seed)
|
||||
mean, std = training_statistics(train_items, seed)
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
manifest = {"train": [str(p) for p, _ in train_items], "val": [str(p) for p, _ in val_items]}
|
||||
(output / "split.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
|
||||
print(f"SPLIT: train_files={len(train_items)} val_files={len(val_items)} overlap=0; mean={mean:.12g} std={std:.12g}", flush=True)
|
||||
active = sorted({label for _, label in train_items})
|
||||
model = GenreCNN(active)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
|
||||
train_loader = torch.utils.data.DataLoader(WavDataset(train_items, mean, std, seed, True), batch_size=batch_size, shuffle=True, num_workers=0)
|
||||
val_loader = torch.utils.data.DataLoader(WavDataset(val_items, mean, std, seed + 1, False), batch_size=batch_size, num_workers=0)
|
||||
best, stale, best_state = math.inf, 0, None
|
||||
checkpoint = output / "genre_model.pt"
|
||||
for number in range(1, epochs + 1):
|
||||
loss, accuracy = epoch(model, train_loader, optimizer)
|
||||
val_loss, val_accuracy = epoch(model, val_loader)
|
||||
print(f"EPOCH {number}/{epochs}: loss={loss:.6f} accuracy={accuracy:.3f} val_loss={val_loss:.6f} val_accuracy={val_accuracy:.3f}", flush=True)
|
||||
if val_loss < best:
|
||||
best, stale = val_loss, 0
|
||||
best_state = copy.deepcopy(model.state_dict())
|
||||
torch.save({"state_dict": best_state, "mean": mean, "std": std, "labels": LABELS, "version": VERSION, "epoch": number}, checkpoint)
|
||||
else:
|
||||
stale += 1
|
||||
if stale >= patience:
|
||||
print(f"EARLY STOP after {number} epochs", flush=True); break
|
||||
model.load_state_dict(best_state); model.eval()
|
||||
model_path, json_path = output / "genre_model.onnx", output / "genre_model.json"
|
||||
example = torch.zeros(1, 1, 128, 300, dtype=torch.float32)
|
||||
# Explicit legacy exporter: opset17, fixed shape, no dynamo/onnxscript requirement.
|
||||
torch.onnx.export(model, example, str(model_path), input_names=["logmel"], output_names=["probabilities"],
|
||||
opset_version=17, dynamo=False, external_data=False)
|
||||
print("EXPORT: torch.onnx.export(dynamo=False), opset=17", flush=True)
|
||||
metadata = {"labels": LABELS, "mean": mean, "std": std, "version": VERSION,
|
||||
"frontend": {"fft": 2048, "hann": "periodic", "hop": "sr//100", "padding": "causal-left-N-hop",
|
||||
"power": "abs(rfft)^2", "mel": "HTK", "bands": 128, "fmin": 30, "fmax": 11000,
|
||||
"area": "divide by full support width in fractional bins", "log_floor": 1e-10}}
|
||||
json_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
|
||||
audio, sr = read_wav(val_items[0][0])
|
||||
validate_export(model_path, model, tensor_frontend(audio, sr, mean, std))
|
||||
print(f"MODEL={model_path.resolve()}\nJSON={json_path.resolve()}\nCHECKPOINT={checkpoint.resolve()}", flush=True)
|
||||
return model_path, json_path
|
||||
|
||||
|
||||
def synthetic_dataset(root: Path, seed: int):
|
||||
rng = np.random.default_rng(seed)
|
||||
sr = 48000
|
||||
time = np.arange(sr * 20, dtype=np.float64) / sr
|
||||
for label_id, label in enumerate(LABELS[:3]):
|
||||
folder = root / label; folder.mkdir(parents=True)
|
||||
for index in range(12):
|
||||
if label_id == 0:
|
||||
period = 60 / rng.uniform(85, 110)
|
||||
phase = np.mod(time, period)
|
||||
audio = np.exp(-phase * 45) * np.sin(2 * np.pi * (70 + index * 3) * time)
|
||||
audio += 0.2 * rng.standard_normal(len(time)) * np.exp(-np.mod(time + period * 0.375, period) * 80)
|
||||
elif label_id == 1:
|
||||
frequency = rng.uniform(150, 450)
|
||||
audio = sum(np.sin(2 * np.pi * frequency * harmonic * time + rng.uniform(0, 6)) / harmonic for harmonic in range(1, 6))
|
||||
audio *= 0.6 + 0.4 * np.sin(2 * np.pi * rng.uniform(0.5, 2) * time)
|
||||
else:
|
||||
audio = lfilter([1, -0.8], [1, -rng.uniform(0.1, 0.6)], rng.standard_normal(len(time)))
|
||||
audio *= 0.4 + 0.6 * np.sin(2 * np.pi * rng.uniform(1, 4) * time) ** 2
|
||||
audio = audio / max(float(np.max(np.abs(audio))), 1e-9) * rng.uniform(0.3, 0.8)
|
||||
wavfile.write(folder / f"synthetic_{index:02d}.wav", sr, np.round(audio * 32767).astype(np.int16))
|
||||
# A fixed, broadband fixture exercises low-energy bins without depending on a player.
|
||||
n = sr * 4
|
||||
t = np.arange(n) / sr
|
||||
fixture = 0.15 * np.sin(2 * np.pi * 440 * t) + 0.1 * np.sin(2 * np.pi * 4500 * t) + 0.03 * rng.standard_normal(n)
|
||||
wavfile.write(root / "fixture.wav", sr, np.round(np.clip(fixture, -1, 1) * 32767).astype(np.int16))
|
||||
print(f"SELFTEST_DIR={root.resolve()}\nFIXTURE={(root / 'fixture.wav').resolve()}", flush=True)
|
||||
|
||||
|
||||
def verify_dump(path: Path, data: Path, rate: int | None, metadata_path: Path | None):
|
||||
if data is None:
|
||||
raise ValueError("--verify-dump needs --data pointing to the printed self-test directory")
|
||||
metadata = json.loads((metadata_path or data / "genre_model.json").read_text(encoding="utf-8"))
|
||||
if metadata["labels"] != LABELS or metadata["version"] != VERSION:
|
||||
raise ValueError("Model metadata contract mismatch")
|
||||
audio, sr = read_wav(data / "fixture.wav")
|
||||
if rate is not None:
|
||||
sr = rate
|
||||
expected = tensor_frontend(audio, sr, metadata["mean"], metadata["std"])
|
||||
if path.stat().st_size != 4 * N_MELS * FRAMES:
|
||||
raise ValueError("Dump must contain exactly 153600 bytes (38400 float32 values)")
|
||||
actual = np.fromfile(path, dtype="<f4").reshape(1, 1, N_MELS, FRAMES)
|
||||
if not np.isfinite(actual).all():
|
||||
raise ValueError("Non-finite tensor dump")
|
||||
difference = float(np.max(np.abs(actual.astype(np.float64) - expected.astype(np.float64))))
|
||||
print(f"PARITY rate={sr} shape={actual.shape} max_abs={difference:.12g} threshold=0.0001", flush=True)
|
||||
if difference >= 1e-4:
|
||||
raise AssertionError(f"Frontend parity failed: {difference}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--data", type=Path)
|
||||
parser.add_argument("--output-dir", type=Path, default=Path(__file__).resolve().parents[1] / "bin")
|
||||
parser.add_argument("--epochs", type=int, default=20)
|
||||
parser.add_argument("--batch-size", type=int, default=4)
|
||||
parser.add_argument("--patience", type=int, default=5)
|
||||
parser.add_argument("--seed", type=int, default=12345)
|
||||
parser.add_argument("--threads", type=int, default=4)
|
||||
parser.add_argument("--selftest", action="store_true")
|
||||
parser.add_argument("--verify-dump", type=Path)
|
||||
parser.add_argument("--metadata-json", type=Path)
|
||||
parser.add_argument("--rate", type=int, choices=[44100, 48000])
|
||||
args = parser.parse_args()
|
||||
if min(args.epochs, args.batch_size, args.patience, args.threads) < 1:
|
||||
parser.error("epochs/batch-size/patience/threads must be positive")
|
||||
torch.set_num_threads(args.threads)
|
||||
if args.verify_dump:
|
||||
verify_dump(args.verify_dump, args.data, args.rate, args.metadata_json)
|
||||
elif args.selftest or args.data is None:
|
||||
root = Path(tempfile.mkdtemp(prefix="ldcc-genre-selftest-"))
|
||||
synthetic_dataset(root, args.seed)
|
||||
train(root, root, 3, args.seed, args.batch_size, args.patience)
|
||||
print("SELFTEST PASS: synthetic only, not a usable musical genre model", flush=True)
|
||||
else:
|
||||
train(args.data, args.output_dir, args.epochs, args.seed, args.batch_size, args.patience)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user