| |
| from pathlib import Path |
| import json, sys, struct |
| import toml, numpy as np, torch, torchaudio |
| import onnx, onnxruntime as ort |
|
|
| TASK_DIR = Path(__file__).resolve().parents[1] |
| EXPORT_DIR = Path(__file__).resolve().parent |
| ORIGIN = TASK_DIR / "origin" |
| _hf = Path((TASK_DIR / "cache/acquire/hf_model_path.txt").read_text().strip()) |
| HF_PATH = _hf if _hf.is_absolute() else (TASK_DIR.parents[2] / _hf).resolve() |
| sys.path.insert(0, str(ORIGIN)) |
| sys.path.insert(0, str(ORIGIN / "pyannote-audio")) |
|
|
| from diarizen.models.eend.model_wavlm_conformer import Model |
|
|
| def main(): |
| torch.manual_seed(0) |
| out_dir = EXPORT_DIR |
| out_dir.mkdir(parents=True, exist_ok=True) |
|
|
| config = toml.load(HF_PATH / "config.toml") |
| args = config["model"]["args"] |
| model = Model(**args) |
| state = torch.load(HF_PATH / "pytorch_model.bin", map_location="cpu") |
| model.load_state_dict(state, strict=True) |
| model.eval() |
|
|
| wav_path = ORIGIN / "example/EN2002a_30s.wav" |
| waveform, sr = torchaudio.load(wav_path) |
| if sr != 16000: |
| waveform = torchaudio.functional.resample(waveform, sr, 16000) |
| waveform = waveform[:1, :4*16000] |
|
|
| full_onnx = onnx.load(TASK_DIR / "export_4s/model.onnx") |
| eps = 1e-5 |
| for n in full_onnx.graph.node: |
| if n.name == '/Constant_2': |
| eps = struct.unpack('f', n.attribute[0].t.raw_data)[0] |
|
|
| wav_1d = waveform[0, :4*16000] |
| mean = wav_1d.mean() |
| var = ((wav_1d - mean) ** 2).mean() |
| wav_ln = (wav_1d - mean) / np.sqrt(var + eps) |
| wav_input = torch.from_numpy(wav_ln).float().unsqueeze(0) |
|
|
| wavlm = model.wavlm_model |
| with torch.no_grad(): |
| cnn_features, _ = wavlm.feature_extractor(wav_input, None) |
| print(f"CNN features shape: {cnn_features.shape}") |
|
|
| class CNNWrapper(torch.nn.Module): |
| def __init__(self, fe): |
| super().__init__() |
| self.fe = fe |
| def forward(self, x): |
| out, _ = self.fe(x, None) |
| return out |
|
|
| cnn_wrapper = CNNWrapper(wavlm.feature_extractor) |
| cnn_wrapper.eval() |
|
|
| onnx_path = out_dir / "cnn_features.onnx" |
| torch.onnx.export( |
| cnn_wrapper, wav_input, onnx_path.as_posix(), |
| input_names=["waveform_ln"], output_names=["cnn_features"], |
| opset_version=16, do_constant_folding=True, dynamic_axes=None, |
| ) |
| onnx_model = onnx.load(onnx_path.as_posix()) |
| onnx.checker.check_model(onnx_model) |
|
|
| sess = ort.InferenceSession(onnx_path.as_posix(), providers=["CPUExecutionProvider"]) |
| ort_out = sess.run(None, {"waveform_ln": wav_input.numpy()})[0] |
| diff = ort_out - cnn_features.numpy() |
| cosine = float(np.dot(ort_out.ravel(), cnn_features.numpy().ravel()) / |
| (np.linalg.norm(ort_out.ravel()) * np.linalg.norm(cnn_features.numpy().ravel()) + 1e-12)) |
| print(f"Cosine: {cosine:.10f}, MAE: {float(np.mean(np.abs(diff))):.8f}") |
|
|
| calib_dir = out_dir / "calib_data" |
| calib_dir.mkdir(exist_ok=True) |
| np.save(calib_dir / "waveform_0.npy", wav_input.numpy().astype(np.float32)) |
|
|
| meta = { |
| "model_name": "diarizen-wavlm-cnn-frontend", |
| "inputs": [{"name": "waveform_ln", "shape": [1, 64000], "dtype": "float32"}], |
| "outputs": [{"name": "cnn_features", "shape": list(ort_out.shape), "dtype": "float32"}], |
| } |
| (out_dir / "model_meta.json").write_text(json.dumps(meta, indent=2) + "\n") |
| print(f"Done: {onnx_path}") |
|
|
| if __name__ == "__main__": |
| main() |
|
|
