File size: 1,261 Bytes
5445376 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | """Compare encoder.onnx against the torch encoder+projector for several audio lengths."""
import numpy as np
import onnxruntime as rt
import soundfile as sf
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
OUT = "/media/hdd16/onnx-asr-exports/granite-speech-3.3-2b"
torch.set_num_threads(8)
torch.set_grad_enabled(False)
processor = AutoProcessor.from_pretrained("ibm-granite/granite-speech-3.3-2b")
model = AutoModelForSpeechSeq2Seq.from_pretrained(
"ibm-granite/granite-speech-3.3-2b", dtype=torch.float32, device_map="cpu"
).eval()
opts = rt.SessionOptions()
opts.intra_op_num_threads = 8
sess = rt.InferenceSession(f"{OUT}/encoder.onnx", opts, providers=["CPUExecutionProvider"])
wav, _ = sf.read(f"{OUT}/clips/en_2.wav", dtype="float32")
for n in (16000, 16000 + 137, 48000, len(wav)):
x = wav[:n][None]
feats = processor.audio_processor(torch.from_numpy(x), device="cpu")
ref = model.projector(model.encoder(feats["input_features"])).numpy()
got = sess.run(["audio_embeds"], {"input_features": x})[0]
exp_len = int(feats["input_features_mask"].sum())
print(
f"n={n:7d} onnx={got.shape} torch={ref.shape} expected_len={exp_len} "
f"max_abs={np.abs(got - ref).max():.3e}"
)
|