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:
@@ -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()
|
||||
Reference in New Issue
Block a user