root
42443420be
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.
59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
#!/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()
|