File size: 6,619 Bytes
38c1604 | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | """
Step 5: export the trained model for the app.
python export.py --config config/pilot.yaml
python export.py --config config/full.yaml --best # mid-training: the best checkpoint so far
Averages the last checkpoints (or takes best.pt with --best), exports ONNX in the interface sherpa-onnx loads as a
"zipformer CTC" model, quantizes it to int8, then checks on test clips that sherpa-onnx (the
engine on the phone) produces the same text as PyTorch. Finally copies model.int8.onnx and
tokens.txt into the app (paths.app_assets).
"""
import copy
import os
import shutil
import sys
import onnx
import torch
from quran_asr import config as config_mod
from quran_asr import data, features, model as model_mod, text
def load_averaged(run_dir, average_last):
steps = sorted(f for f in os.listdir(run_dir) if f.startswith("step_") and f.endswith(".pt"))
chosen = steps[-average_last:] if steps and average_last > 1 else []
if not chosen:
chosen = ["best.pt"] if os.path.exists(os.path.join(run_dir, "best.pt")) else steps[-1:]
print("averaging:", ", ".join(chosen))
states = [torch.load(os.path.join(run_dir, f), map_location="cpu", weights_only=False) for f in chosen]
avg = {}
for key, value in states[0]["model"].items():
if value.is_floating_point():
avg[key] = sum(s["model"][key] for s in states) / len(states)
else:
avg[key] = value
return avg, states[0]["model_config"]
class Exportable(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, x, x_lens):
log_probs, lens = self.model(x, x_lens)
return log_probs, lens.to(torch.int64)
def export_onnx(model, path):
model.eval()
x = torch.randn(1, 600, features.NUM_BINS)
x_lens = torch.tensor([600], dtype=torch.int64)
torch.onnx.export(
Exportable(model), (x, x_lens), path, opset_version=17, dynamo=False,
input_names=["x", "x_lens"], output_names=["log_probs", "log_probs_len"],
dynamic_axes={"x": {0: "N", 1: "T"}, "x_lens": {0: "N"}, "log_probs": {0: "N", 1: "T_out"},
"log_probs_len": {0: "N"}},
)
m = onnx.load(path)
meta = {"model_type": "zipformer2_ctc", "vocab_size": str(len(text.VOCAB)), "subsampling_factor": "4",
"comment": "Quran recitation, Arabic letters without harakat (Warsh + Hafs)"}
for k, v in meta.items():
entry = m.metadata_props.add()
entry.key, entry.value = k, v
onnx.save(m, path)
def quantize(src, dst):
from onnxruntime.quantization import QuantType, quantize_dynamic
quantize_dynamic(src, dst, op_types_to_quantize=["MatMul"], weight_type=QuantType.QInt8)
def sherpa_recognizer(model_path, tokens_path):
import sherpa_onnx
return sherpa_onnx.OfflineRecognizer.from_zipformer_ctc(model=model_path, tokens=tokens_path, num_threads=4)
def sherpa_text(recognizer, wave):
stream = recognizer.create_stream()
stream.accept_waveform(features.SAMPLE_RATE, wave)
recognizer.decode_stream(stream)
return stream.result.text.replace(text.WORD, " ").strip()
def main():
parser = config_mod.from_args(__doc__)
parser.add_argument("--best", action="store_true", help="export best.pt (lowest validation error so far)")
args = parser.parse_args()
cfg = config_mod.load(args.config)
torch.set_num_threads(4) # leave the rest of the CPU to a running training
run = cfg.paths.run_dir
if args.best:
ckpt = torch.load(os.path.join(run, "best.pt"), map_location="cpu", weights_only=False)
print(f"best.pt: step {ckpt.get('step')}, validation letter error {ckpt.get('valid_cer', 0) * 100:.2f}%")
state, model_config = ckpt["model"], ckpt["model_config"]
else:
state, model_config = load_averaged(run, cfg.export.average_last)
model = model_mod.build(config_mod._wrap(model_config), len(text.VOCAB))
model.load_state_dict(state)
model_mod.use_plain_attention(model).eval()
out = os.path.join(run, "export")
os.makedirs(out, exist_ok=True)
fp32 = os.path.join(out, "model.onnx")
export_onnx(copy.deepcopy(model), fp32) # tracing changes the model it exports: keep ours clean for the parity check
final = fp32
if cfg.export.quantize == "int8":
final = os.path.join(out, "model.int8.onnx")
quantize(fp32, final)
tokens = os.path.join(out, "tokens.txt")
text.write_tokens(tokens)
print(f"exported {final} ({os.path.getsize(final) / 1e6:.1f} MB)")
# Parity: the phone's engine must hear exactly what PyTorch hears.
from quran_asr import audio
test = data.read_manifest(os.path.join(config_mod.manifest_dir(cfg), "test.jsonl"))
clips = test[:: max(1, len(test) // cfg.export.parity_clips)][: cfg.export.parity_clips]
rec_final = sherpa_recognizer(final, tokens)
rec_fp32 = sherpa_recognizer(fp32, tokens)
same = 0
errs = {"pytorch": [0, 0], "phone": [0, 0]}
for e in clips:
wave = audio.load(e["path"])
feats = features.fbank(torch.from_numpy(wave))
with torch.no_grad():
log_probs, lens = model(feats[None], torch.tensor([feats.shape[0]]))
ref_text = text.decode(model_mod.greedy_decode(log_probs, lens)[0])
phone_text = sherpa_text(rec_final, wave)
same += ref_text == sherpa_text(rec_fp32, wave) # the export itself must be exact
for name, hyp in (("pytorch", ref_text), ("phone", phone_text)):
ce, ct, *_ = text.error_counts(hyp, e["text"])
errs[name][0] += ce
errs[name][1] += ct
cer = {k: v[0] / max(1, v[1]) for k, v in errs.items()}
print(f"parity on {len(clips)} test clips: sherpa-onnx fp32 = PyTorch on {same}/{len(clips)}; letter error "
f"PyTorch fp32 {cer['pytorch'] * 100:.2f}% vs sherpa-onnx {cfg.export.quantize} {cer['phone'] * 100:.2f}%")
if same < 0.9 * len(clips):
print("WARNING: sherpa-onnx and PyTorch disagree: the export is broken; not copying it into the app")
return 1
if cer["phone"] > cer["pytorch"] + 0.02:
print("WARNING: the exported model is noticeably worse than PyTorch; not copying it into the app")
return 1
os.makedirs(cfg.paths.app_assets, exist_ok=True)
shutil.copy(final, os.path.join(cfg.paths.app_assets, "model.int8.onnx"))
shutil.copy(tokens, os.path.join(cfg.paths.app_assets, "tokens.txt"))
print("copied into", cfg.paths.app_assets, "- rebuild the app to ship it")
return 0
if __name__ == "__main__":
sys.exit(main())
|