chore: OpenVINO MetricGAN GPU spike (do not merge)

Arc B580 lost the kill test: batch-14 mask p50 103 ms vs Torch CPU
70 ms enhance / OpenVINO CPU 5 ms. Production stays on known-good
CPU workers (tag known-good-cpu-2026-08-28). Hop stays 200 ms.
This commit is contained in:
root
2026-08-28 07:18:54 +00:00
parent 736eccb932
commit 42443420be
6 changed files with 307 additions and 0 deletions
+58
View File
@@ -0,0 +1,58 @@
#!/usr/bin/env python3
"""Export MetricGAN enhance_model (BLSTM masker) to ONNX. Uses project .venv."""
from __future__ import annotations
import argparse
from pathlib import Path
import torch
import torch.nn as nn
from speechbrain.inference.enhancement import SpectralMaskEnhancement
SAVEDIR = "/root/.cache/speechbrain-enhancement"
class MaskNet(nn.Module):
def __init__(self, net: nn.Module):
super().__init__()
self.net = net
def forward(self, feats: torch.Tensor) -> torch.Tensor:
# lengths=1.0 relative, one per batch row
lengths = torch.ones(feats.shape[0], dtype=feats.dtype)
return self.net(feats, lengths=lengths)
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--out", default="scripts/metricgan_enhance.onnx")
args = ap.parse_args()
out = Path(args.out)
out.parent.mkdir(parents=True, exist_ok=True)
m = SpectralMaskEnhancement.from_hparams(
source="speechbrain/metricgan-plus-voicebank",
savedir=SAVEDIR,
run_opts={"device": "cpu"},
)
wrap = MaskNet(m.mods.enhance_model).eval()
feats = m.compute_features(torch.randn(1, 16000))
with torch.no_grad():
torch.onnx.export(
wrap,
feats,
str(out),
input_names=["feats"],
output_names=["mask"],
dynamic_axes={
"feats": {0: "batch", 1: "time"},
"mask": {0: "batch", 1: "time"},
},
opset_version=17,
dynamo=False,
)
print("wrote", out, "bytes", out.stat().st_size, "example_feats", tuple(feats.shape))
if __name__ == "__main__":
main()