| """The fused attention used in training and the plain formula used for ONNX must agree.""" |
| import os |
| import sys |
|
|
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| from quran_asr import model as model_mod, text |
| from quran_asr.config import _wrap |
|
|
|
|
| def main(): |
| torch.manual_seed(0) |
| cfg = _wrap(dict(d_model=96, num_layers=3, num_heads=4, ff_mult=2, conv_kernel=15, dropout=0.0, |
| subsampling_channels=64)) |
| m = model_mod.build(cfg, len(text.VOCAB)).eval() |
| x = torch.randn(3, 700, 80) |
| lens = torch.tensor([700, 450, 123]) |
| with torch.no_grad(): |
| fused, l1 = m(x, lens) |
| plain, l2 = model_mod.use_plain_attention(m)(x, lens) |
| keep = torch.arange(fused.shape[1])[None, :] < l1[:, None] |
| diff = (fused - plain).abs()[keep].max().item() |
| print(f"max |fused - plain| over valid frames: {diff:.2e}") |
| assert torch.equal(l1, l2) and diff < 1e-4 |
| print("OK: training and export attention agree") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
