362 lines
17 KiB
Python
362 lines
17 KiB
Python
"""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()
|