export.py

import argparse
import torch
import numpy as np
import onnx
import onnxruntime as ort

from model import SpeechEnhancer


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--checkpoint', default='checkpoints/best.tar')
    p.add_argument('--output',     default='checkpoints/model.onnx')
    p.add_argument('--opset',      type=int, default=18)
    args = p.parse_args()

    device = torch.device('cpu')
    model = SpeechEnhancer().to(device)
    ckpt = torch.load(args.checkpoint, map_location=device)
    model.load_state_dict(ckpt['model'])
    model.eval()

    # Dummy inputs matching the model's forward signature. T is arbitrary here;
    # the dynamic axis below lets the exported model accept any length.
    T = 100
    noisy_feat = torch.randn(1, 3, 257, T)
    noisy_spec = torch.randn(1, 257, T, 2)

    torch.onnx.export(
        model,
        (noisy_feat, noisy_spec),
        args.output,
        input_names=['noisy_feat', 'noisy_spec'],
        output_names=['enhanced_spec'],
        dynamic_axes={
            'noisy_feat':    {3: 'time'},
            'noisy_spec':    {2: 'time'},
            'enhanced_spec': {2: 'time'},
        },
        opset_version=args.opset,
        dynamo=False,
    )

    onnx.checker.check_model(onnx.load(args.output))
    print(f"Exported to {args.output}")

    # Verify ONNX output matches PyTorch within tolerance.
    with torch.no_grad():
        torch_out = model(noisy_feat, noisy_spec).numpy()

    sess = ort.InferenceSession(args.output, providers=['CPUExecutionProvider'])
    onnx_out = sess.run(None, {
        'noisy_feat': noisy_feat.numpy(),
        'noisy_spec': noisy_spec.numpy(),
    })[0]

    max_diff = np.abs(torch_out - onnx_out).max()
    print(f"Max abs diff (PyTorch vs ONNX): {max_diff:.2e}")
    # Recurrent layers accumulate small floating-point differences across
    # timesteps, so a diff on the order of 1e-3 is expected here rather than 1e-6.
    if max_diff < 5e-3:
        print("Export verified.")
    else:
        print("WARNING: outputs differ more than expected.")

    # Confirm a different length works through the dynamic axis.
    T2 = 250
    onnx_out2 = sess.run(None, {
        'noisy_feat': np.random.randn(1, 3, 257, T2).astype(np.float32),
        'noisy_spec': np.random.randn(1, 257, T2, 2).astype(np.float32),
    })[0]
    print(f"Variable-length check: input T={T2} -> output shape {onnx_out2.shape}")


if __name__ == '__main__':
    main()