vision_encoder.onnx has the 224x224 attention scale baked in

#11
by georgeguimaraes - opened

Hey! Looks like vision_encoder.onnx was traced at 224x224, so the DaViT channel attention scale (tokens ** -0.5) got baked in as constants for that size: the 12 .../channel_attn/fn/Mul nodes have 1/56, 1/28, 1/14 and 1/7. The processor resizes to 768x768, where they should be 1/192, 1/96, 1/48 and 1/24, so each of those softmaxes ends up ~3.4x too sharp.

Captions mostly survive, but the text drifts on every task. On a handwriting sample, <OCR_WITH_REGION> gives "Meet the at the station at noon." where PyTorch gives "Meet me at the station 'at noon.". With the 12 constants patched to the 768 values, the export matches PyTorch on all 6 tasks I tried (features within 2e-4).

Re-exporting at 768 (or keeping that scale dynamic) should fix it. Thanks for the exports!

Sign up or log in to comment