7b92c9cAdd dataset generation scripts, ML fitting code and reference recordingsJeremy Magland 1"""Trains a convolutional network that predicts the sheet parameters from the
2spectrogram of a strike, evaluates it on held-out clips, and writes a report.
4 python3 ml/train.py --data DATASET_DIR --out RUN_DIR [--cache CACHE_DIR] [--epochs N] [--val-every K] [--device cuda|cpu]
6Targets are the logarithms of E, rho, h, Lx, Ly, standardized by the training
7set; the loss is the mean squared error over the five. Every K-th clip by
8index is held out for validation. The report gives, per target and for a few
9derived quantities (E/rho, the bending wave constant kappa, the area, the
10aspect ratio), the relative RMS error of the prediction against that of
11predicting the training mean, and scatter plots. A line is appended to
12RUN_DIR/../history.jsonl so that the error can be tracked as clips accumulate.
13"""
15import argparse
16import base64
17import io
18import json
19import math
20import os
21import sys
22import time
24import numpy as np
25import torch
26import torch.nn as nn
28# The office machine's PyTorch build cannot find a cuDNN engine for the
29# convolution backward pass; the native kernels work, more slowly.
30torch.backends.cudnn.enabled = os.environ.get("CUDNN", "0") == "1"
32sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
33from features import build_cache # noqa: E402
35TARGETS = ["E", "rho", "h", "Lx", "Ly"]
36UNITS = {"E": ("GPa", 1e-9), "rho": ("kg/m^3", 1), "h": ("mm", 1e3), "Lx": ("in", 1 / 0.0254), "Ly": ("in", 1 / 0.0254)}
37DERIVED = {
38 "E/rho": lambda d: d["E"] - d["rho"],
39 "kappa": lambda d: d["h"] + 0.5 * (d["E"] - d["rho"]),
40 "area": lambda d: d["Lx"] + d["Ly"],
41 "aspect": lambda d: d["Lx"] - d["Ly"],
42 "kappa/area": lambda d: d["h"] + 0.5 * (d["E"] - d["rho"]) - d["Lx"] - d["Ly"],
43}
46class Net(nn.Module):
47 """Six strided conv blocks, global average pooling, and a small head. A
48 second input channel carries the frequency coordinate of each row, so the
49 network can read absolute frequency despite the pooling."""
51 def __init__(self, n_out=5, channels=(16, 32, 64, 96, 128, 128)):
52 super().__init__()
53 layers, c_in = [], 2
54 for c in channels:
55 layers += [nn.Conv2d(c_in, c, 3, padding=1), nn.BatchNorm2d(c), nn.ReLU(inplace=True),
56 nn.Conv2d(c, c, 3, padding=1), nn.BatchNorm2d(c), nn.ReLU(inplace=True), nn.MaxPool2d(2)]
57 c_in = c
58 self.features = nn.Sequential(*layers)
59 self.head = nn.Sequential(nn.Linear(c_in, 128), nn.ReLU(inplace=True), nn.Dropout(0.2), nn.Linear(128, n_out))
61 def forward(self, x):
62 b, _, f, t = x.shape
63 coord = torch.linspace(0, 1, f, device=x.device).view(1, 1, f, 1).expand(b, 1, f, t)
64 h = self.features(torch.cat([x, coord], 1)).mean((2, 3))
65 return self.head(h)
68class SpecNet(nn.Module):
69 """Full-frequency-resolution model. The spectrogram is collapsed into a few
70 time slices (mean log spectrum over 0-0.25, 0.25-1, 1-2, 2-3 s), which keep
71 the partials at the resolution of the bins and the decay of each band as
72 the ratio between slices. A 1D convolutional stack along frequency (two
73 stride-2 steps only, dilations for context) is followed by multi-head
74 attention pooling over frequency positions with a sinusoidal position code,
75 so that the pooled vector can express where the lines are."""
77 def __init__(self, n_out=5, n_freq=1024, fps=48000 / 1024, slices=((0, 0.25), (0.25, 1), (1, 2), (2, 3)), heads=4, pos_dim=16):
78 super().__init__()
79 self.bounds = [(int(a * fps), int(b * fps)) for a, b in slices]
80 c_in = len(slices) + 1
81 def block(ci, co, k, stride=1, dil=1):
82 return [nn.Conv1d(ci, co, k, stride=stride, padding=dil * (k // 2), dilation=dil), nn.BatchNorm1d(co), nn.ReLU(inplace=True)]
83 self.conv = nn.Sequential(*block(c_in, 32, 7), *block(32, 32, 7), *block(32, 64, 5, stride=2), *block(64, 64, 5, dil=2),
84 *block(64, 96, 5, stride=2), *block(96, 96, 5, dil=4), *block(96, 96, 5, dil=8), *block(96, 128, 3))
85 n_pos = (n_freq + 3) // 4
86 pos = torch.arange(n_pos, dtype=torch.float32)[:, None] / n_pos
87 freqs = 2 ** torch.arange(pos_dim // 2, dtype=torch.float32)
88 self.register_buffer("pos", torch.cat([torch.sin(math.pi * pos * freqs), torch.cos(math.pi * pos * freqs)], 1))
89 d = 128 + pos_dim
90 self.att = nn.Linear(d, heads)
91 self.head = nn.Sequential(nn.Linear(heads * d, 256), nn.ReLU(inplace=True), nn.Dropout(0.2), nn.Linear(256, n_out))
93 def forward(self, x):
94 b, _, f, t = x.shape
95 s = torch.stack([x[:, 0, :, a:min(b_, t)].mean(-1) for a, b_ in self.bounds], 1)
96 coord = torch.linspace(0, 1, f, device=x.device).view(1, 1, f).expand(b, 1, f)
97 h = self.conv(torch.cat([s, coord], 1)).transpose(1, 2) # (b, positions, 128)
98 hp = torch.cat([h, self.pos[: h.shape[1]].expand(b, -1, -1)], -1)
99 a = torch.softmax(self.att(hp), dim=1) # (b, positions, heads)
100 pooled = torch.einsum("bpk,bpd->bkd", a, hp).flatten(1)
101 return self.head(pooled)
104# Alternative target basis: the combinations the linear physics identifies
105# (kappa / area up to a constant, aspect, E / rho, area) plus rho, so that the
106# well-determined directions get their own outputs and clean gradients.
107COMBOS = np.array([
108 [0.5, -0.5, 1, -1, -1], # log kappa/area (up to a constant)
109 [0, 0, 0, 1, -1], # log aspect
110 [1, -1, 0, 0, 0], # log E/rho
111 [0, 0, 0, 1, 1], # log area
112 [0, 1, 0, 0, 0], # log rho
113], dtype=np.float32)
116def metrics(pred, true, mean):
117 """Relative RMS error (percent) of pred and of the constant mean, and R^2, all in log space."""
118 err = np.sqrt(np.mean((pred - true) ** 2))
119 base = np.sqrt(np.mean((mean - true) ** 2))
120 return {"rel_rms_pct": float(100 * err), "baseline_rel_rms_pct": float(100 * base), "r2": float(1 - err ** 2 / max(base ** 2, 1e-12)), "mae_pct": float(100 * np.mean(np.abs(pred - true)))}
123def scatter_png(true, pred, name, unit, scale):
124 import matplotlib
125 matplotlib.use("Agg")
126 import matplotlib.pyplot as plt
127 fig, ax = plt.subplots(figsize=(3.2, 3.2), dpi=100)
128 t, p = np.exp(true) * scale, np.exp(pred) * scale
129 ax.scatter(t, p, s=6, alpha=0.5)
130 lo, hi = min(t.min(), p.min()), max(t.max(), p.max())
131 ax.plot([lo, hi], [lo, hi], "k-", lw=0.8)
132 ax.set_xlabel(f"true {name} [{unit}]")
133 ax.set_ylabel(f"predicted {name} [{unit}]")
134 ax.set_xscale("log"); ax.set_yscale("log")
135 fig.tight_layout()
136 buf = io.BytesIO(); fig.savefig(buf, format="png"); plt.close(fig)
137 return base64.b64encode(buf.getvalue()).decode()
140def main():
141 ap = argparse.ArgumentParser()
142 ap.add_argument("--data", required=True)
143 ap.add_argument("--out", required=True)
144 ap.add_argument("--cache", default=None)
145 ap.add_argument("--epochs", type=int, default=60)
146 ap.add_argument("--batch", type=int, default=32)
147 ap.add_argument("--lr", type=float, default=1e-3)
148 ap.add_argument("--val-every", type=int, default=8, help="every k-th clip by index is held out")
149 ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
150 ap.add_argument("--workers", type=int, default=8)
151 ap.add_argument("--logf", action="store_true", help="log-frequency spectrogram input")
152 ap.add_argument("--no-augment", action="store_true", help="disable the onset-shift and masking augmentation")
153 ap.add_argument("--arch", default="cnn2d", choices=["cnn2d", "spec1d"], help="pooling 2D CNN (default) or the full-resolution SpecNet")
154 ap.add_argument("--targets", default="params", choices=["params", "combos"], help="train on log E, rho, h, Lx, Ly or on the identifiable combinations")
155 ap.add_argument("--tag", default=None, help="suffix of the report and history files (default: 'logf' with --logf, else none)")
156 args = ap.parse_args()
157 os.makedirs(args.out, exist_ok=True)
158 kind = "logf" if args.logf else "lin"
159 tag = args.tag if args.tag is not None else "-".join(t for t in [("logf" if args.logf else ""), ("" if args.arch == "cnn2d" else args.arch), ("" if args.targets == "params" else args.targets)] if t)
160 suffix = f"-{tag}" if tag else ""
161 cache = args.cache or os.path.join(args.data, "ml", "cache" + ("-logf" if args.logf else ""))
162 t_start = time.time()
163 rows, feats = build_cache(args.data, cache, args.workers, kind)
164 n = len(rows)
165 y = np.stack([[math.log(r[t]) for t in TARGETS] for r in rows]).astype(np.float32)
166 idx = np.array([r["index"] for r in rows])
167 val = idx % args.val_every == 0
168 tr = ~val
169 if tr.sum() < 20 or val.sum() < 5:
170 print(f"only {n} clips; need more"); return
171 T = COMBOS if args.targets == "combos" else np.eye(len(TARGETS), dtype=np.float32)
172 Tinv = np.linalg.inv(T)
173 yt = y @ T.T
174 mu, sd = yt[tr].mean(0), yt[tr].std(0) + 1e-6
175 z = (yt - mu) / sd
176 print(f"{n} clips: {tr.sum()} train, {val.sum()} val; features {feats.shape}; device {args.device}")
178 dev = torch.device(args.device)
179 X = torch.from_numpy(feats) # float16 (n, F, T), kept on the CPU
180 Z = torch.from_numpy(z)
181 net = (SpecNet(len(TARGETS), n_freq=feats.shape[1]) if args.arch == "spec1d" else Net(len(TARGETS))).to(dev)
182 print(f"arch {args.arch}, targets {args.targets}, {sum(p.numel() for p in net.parameters()) / 1e6:.2f} M parameters")
183 opt = torch.optim.AdamW(net.parameters(), lr=args.lr, weight_decay=1e-4)
184 sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=args.epochs * math.ceil(tr.sum() / args.batch))
185 tr_idx, val_idx = np.where(tr)[0], np.where(val)[0]
186 rng = np.random.default_rng(0)
188 def predict(ids):
189 net.eval()
190 out = []
191 with torch.no_grad():
192 for k in range(0, len(ids), 64):
193 xb = X[ids[k:k + 64]].to(dev).float().unsqueeze(1)
194 out.append(net(xb).cpu().numpy())
195 return np.concatenate(out)
197 def augment(xb):
198 """Label-preserving perturbations against memorization: a random onset
199 shift of up to 8 frames (170 ms), and one masked frequency band and one
200 masked time span per clip (the masked cells take the floor value 0)."""
201 bsz, _, nf, nt = xb.shape
202 out = torch.zeros_like(xb)
203 for i in range(bsz):
204 sh = int(rng.integers(0, 9))
205 out[i, :, :, sh:] = xb[i, :, :, : nt - sh]
206 f0, fw = int(rng.integers(0, nf)), int(rng.integers(0, nf // 16))
207 out[i, :, f0:f0 + fw, :] = 0
208 t0, tw = int(rng.integers(0, nt)), int(rng.integers(0, nt // 8))
209 out[i, :, :, t0:t0 + tw] = 0
210 return out
212 history, best, best_state = [], float("inf"), None
213 for epoch in range(args.epochs):
214 net.train()
215 perm = rng.permutation(tr_idx)
216 tot = 0.0
217 for k in range(0, len(perm), args.batch):
218 b = perm[k:k + args.batch]
219 xb = X[b].to(dev).float().unsqueeze(1)
220 if not args.no_augment:
221 xb = augment(xb)
222 zb = Z[b].to(dev)
223 loss = nn.functional.mse_loss(net(xb), zb)
224 opt.zero_grad(); loss.backward(); opt.step(); sched.step()
225 tot += loss.item() * len(b)
226 pv = predict(val_idx)
227 vloss = float(np.mean((pv - z[val_idx]) ** 2))
228 history.append({"epoch": epoch, "train_loss": tot / len(perm), "val_loss": vloss})
229 if vloss < best:
230 best, best_state = vloss, {k: v.detach().cpu().clone() for k, v in net.state_dict().items()}
231 print(f"epoch {epoch}: train {tot / len(perm):.4f} val {vloss:.4f}", flush=True)
232 net.load_state_dict(best_state)
233 torch.save({"state": best_state, "mu": mu, "sd": sd, "targets": TARGETS, "T": T, "arch": args.arch, "input": kind, "n_freq": feats.shape[1]}, os.path.join(args.out, "model.pt"))
235 # Evaluation in log space, back in the natural scale of each target.
236 pv = (predict(val_idx) * sd + mu) @ Tinv.T
237 yv = y[val_idx]
238 ym = np.broadcast_to(mu @ Tinv.T, yv.shape)
239 res = {t: metrics(pv[:, i], yv[:, i], ym[:, i]) for i, t in enumerate(TARGETS)}
240 d_pred = {t: pv[:, i] for i, t in enumerate(TARGETS)}
241 d_true = {t: yv[:, i] for i, t in enumerate(TARGETS)}
242 d_mean = {t: ym[:, i] for i, t in enumerate(TARGETS)}
243 for name, f in DERIVED.items():
244 res[name] = metrics(f(d_pred), f(d_true), f(d_mean))
245 summary = {"input": kind, "arch": args.arch, "targets": args.targets, "augment": not args.no_augment, "n_clips": n, "n_train": int(tr.sum()), "n_val": int(val.sum()), "best_val_loss": float(best), "epochs": args.epochs,
246 "minutes": (time.time() - t_start) / 60, "finished": time.strftime("%Y-%m-%d %H:%M:%S"), "metrics": res, "history": history}
247 with open(os.path.join(args.out, "metrics.json"), "w") as f:
248 json.dump(summary, f, indent=1)
249 with open(os.path.join(args.out, "..", f"history{suffix}.jsonl"), "a") as f:
250 f.write(json.dumps({k: v for k, v in summary.items() if k != "history"}) + "\n")
251 np.save(os.path.join(args.out, "val_pred.npy"), np.stack([idx[val_idx], *yv.T, *pv.T], 1))
252 write_report(args.out, summary, yv, pv, suffix)
253 print("\n".join(f"{t:>8}: rel RMS {m['rel_rms_pct']:.1f}% (mean-baseline {m['baseline_rel_rms_pct']:.1f}%), R^2 {m['r2']:.2f}" for t, m in res.items()))
256def write_report(out, summary, yv, pv, suffix=""):
257 hist_lines = []
258 hp = os.path.join(out, "..", f"history{suffix}.jsonl")
259 if os.path.exists(hp):
260 with open(hp) as f:
261 hist_lines = [json.loads(l) for l in f if l.strip()]
262 rows = "".join(
263 f"<tr><td>{t}</td><td>{m['rel_rms_pct']:.1f}</td><td>{m['mae_pct']:.1f}</td><td>{m['baseline_rel_rms_pct']:.1f}</td><td>{m['r2']:.2f}</td></tr>"
264 for t, m in summary["metrics"].items())
265 imgs = "".join(f'<img src="data:image/png;base64,{scatter_png(yv[:, i], pv[:, i], t, *UNITS[t])}">' for i, t in enumerate(TARGETS))
266 hist_rows = "".join(f"<tr><td>{h['finished']}</td><td>{h['n_clips']}</td>" + "".join(f"<td>{h['metrics'][t]['rel_rms_pct']:.1f}</td>" for t in TARGETS) + "</tr>" for h in hist_lines)
267 curve = "".join(f"<tr><td>{h['epoch']}</td><td>{h['train_loss']:.4f}</td><td>{h['val_loss']:.4f}</td></tr>" for h in summary["history"][::max(1, len(summary["history"]) // 15)])
268 html = f"""<!doctype html><meta charset="utf-8"><title>steelsim31 prediction report</title>
269<style>body{{font:14px system-ui,sans-serif;margin:16px;color:#222}} table{{border-collapse:collapse;margin:8px 0}} td,th{{padding:3px 10px;border-bottom:1px solid #ddd;text-align:right}} td:first-child,th:first-child{{text-align:left}} img{{margin:4px}} h2{{font-size:16px;margin-top:20px}}</style>
270<h1 style="font-size:18px">Prediction of the sheet parameters from the strike sound</h1>
271<p>Input: {"log-frequency" if summary["input"] == "logf" else "linear-frequency"} spectrogram; architecture {summary.get("arch", "cnn2d")}; targets {summary.get("targets", "params")}. {summary['n_clips']} clips ({summary['n_train']} train, {summary['n_val']} held out), {summary['epochs']} epochs, {summary['minutes']:.1f} min, finished {summary['finished']}. Errors are relative (percent, from the logarithm) on the held-out clips; the baseline predicts the training mean.</p>
272<table><tr><th>target</th><th>rel RMS %</th><th>rel MAE %</th><th>baseline rel RMS %</th><th>R²</th></tr>{rows}</table>
273<p>Derived: E/rho (sets the wave speeds), kappa = h sqrt(E / (12 rho (1 - nu²))) (sets the bending frequencies), area = Lx Ly, aspect = Lx / Ly, kappa/area (the scale of the bending spectrum).</p>
274<h2>Held-out predictions</h2>{imgs}
275<h2>Error against the number of clips</h2>
276<table><tr><th>run</th><th>clips</th>{"".join(f"<th>{t} %</th>" for t in TARGETS)}</tr>{hist_rows}</table>
277<h2>Training curve (standardized MSE)</h2>
278<table><tr><th>epoch</th><th>train</th><th>val</th></tr>{curve}</table>
279"""
280 with open(os.path.join(out, "report.html"), "w") as f:
281 f.write(html)
282 with open(os.path.join(out, "..", f"report{suffix}.html"), "w") as f:
283 f.write(html)
286if __name__ == "__main__":
287 main()