Upload ONNX export

export_pp_doclayout_v3.py

@@ -0,0 +1,215 @@
+"""Export PaddlePaddle/PP-DocLayoutV3_safetensors to ONNX.
+
+The HF module returns a large structured output (intermediate decoder states,
+per-layer masks, reference points...). Only four tensors are needed at serving
+time, so we wrap the model and return exactly those:
+
+    logits        (B, 300, num_classes)
+    pred_boxes    (B, 300, 4)          cxcywh, normalised
+    order_logits  (B, 300, 300)        reading-order pointer matrix
+    out_masks     (B, 300, 200, 200)   mask logits @ stride 4  (optional)
+
+`out_masks` is by far the biggest tensor (300*200*200*4B = 48 MB per image in
+fp32). Export with --no-masks if you only need boxes + reading order; polygons
+then degrade to axis-aligned rectangles.
+
+Usage:
+    python export_pp_doclayout_v3.py --output pp_doclayoutv3.onnx
+    python export_pp_doclayout_v3.py --output pp_doclayoutv3_nomask.onnx --no-masks
+    python export_pp_doclayout_v3.py --output m.onnx --fp16   # extra fp16 copy
+"""
+
+from __future__ import annotations
+
+import argparse
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import nn
+from transformers import AutoModelForObjectDetection
+
+DEFAULT_MODEL = "PaddlePaddle/PP-DocLayoutV3_safetensors"
+INPUT_SIZE = 800  # preprocessor_config.json: {"height": 800, "width": 800}
+
+
+@contextmanager
+def fp32_position_embedding():
+    """Force the 2D sin/cos position embedding to float32 during tracing.
+
+    Upstream computes the frequency grid in float64, which bakes Sin/Cos(double)
+    nodes into the graph. ONNX Runtime's CPU EP has no double kernel for those, so
+    the session fails to load with:
+        NOT_IMPLEMENTED : Could not find an implementation for Cos(7)
+    The embedding is cast to float32 anyway; the difference is ~1e-6.
+    """
+    from transformers.models.pp_doclayout_v3 import modeling_pp_doclayout_v3 as M
+
+    original = M.build_2d_sinusoidal_position_embedding
+
+    def patched(height, width, embed_dim=256, temperature=10000.0,
+                cls_token=False, device=None, dtype=torch.float32):
+        pos_dim = embed_dim // 4
+        omega = torch.arange(pos_dim, dtype=torch.float32, device=device) / pos_dim
+        omega = 1.0 / temperature**omega
+        grid_h = torch.arange(height, dtype=torch.float32, device=device)
+        grid_w = torch.arange(width, dtype=torch.float32, device=device)
+        grid_h, grid_w = torch.meshgrid(grid_h, grid_w, indexing="ij")
+        emb_h = grid_h.flatten().outer(omega)
+        emb_w = grid_w.flatten().outer(omega)
+        pos_embed = torch.cat([emb_h.sin(), emb_h.cos(), emb_w.sin(), emb_w.cos()], dim=1)
+        if cls_token:
+            zeros = torch.zeros(1, embed_dim, dtype=torch.float32, device=device)
+            pos_embed = torch.cat([zeros, pos_embed], dim=0)
+        return pos_embed.to(dtype)
+
+    M.build_2d_sinusoidal_position_embedding = patched
+    try:
+        yield
+    finally:
+        M.build_2d_sinusoidal_position_embedding = original
+
+
+class PPDocLayoutV3ExportWrapper(nn.Module):
+    """Flattens the HF output struct down to the tensors post-processing needs."""
+
+    def __init__(self, model: nn.Module, with_masks: bool = True) -> None:
+        super().__init__()
+        self.model = model
+        self.with_masks = with_masks
+
+    def forward(self, pixel_values: torch.Tensor):
+        out = self.model(pixel_values=pixel_values)
+
+        order_logits = out.order_logits
+        # Doc types it as a tuple; take the final decoder layer if so.
+        if isinstance(order_logits, (tuple, list)):
+            order_logits = order_logits[-1]
+        if order_logits.dim() == 4:  # (B, layers, Q, Q)
+            order_logits = order_logits[:, -1]
+
+        if not self.with_masks:
+            return out.logits, out.pred_boxes, order_logits
+
+        masks = out.out_masks
+        if masks.dim() == 5:  # (B, layers, Q, H, W) -> last layer only
+            masks = masks[:, -1]
+        return out.logits, out.pred_boxes, order_logits, masks
+
+
+def export(
+    model_path: str,
+    output: Path,
+    with_masks: bool = True,
+    opset: int = 17,
+    dynamo: bool = False,
+    dynamic_batch: bool = True,
+) -> tuple[nn.Module, list[str]]:
+    model = AutoModelForObjectDetection.from_pretrained(model_path, dtype=torch.float32).eval()
+    # Pure-PyTorch deformable attention (grid_sample) instead of the custom CUDA
+    # kernel — the custom op has no ONNX symbolic. The checkpoint config already
+    # sets this, but be explicit in case someone overrides it.
+    model.config.disable_custom_kernels = True
+
+    wrapper = PPDocLayoutV3ExportWrapper(model, with_masks=with_masks).eval()
+
+    dummy = torch.randn(1, 3, INPUT_SIZE, INPUT_SIZE)
+    output_names = ["logits", "pred_boxes", "order_logits"]
+    if with_masks:
+        output_names.append("out_masks")
+
+    # Height/width stay static: the image processor always resizes to 800x800,
+    # and static spatial dims let ORT constant-fold the anchor generation.
+    dynamic_axes = None
+    if dynamic_batch:
+        dynamic_axes = {name: {0: "batch"} for name in ["pixel_values"] + output_names}
+
+    output.parent.mkdir(parents=True, exist_ok=True)
+    with torch.inference_mode(), fp32_position_embedding():
+        torch.onnx.export(
+            wrapper,
+            (dummy,),
+            str(output),
+            input_names=["pixel_values"],
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
+            opset_version=opset,
+            do_constant_folding=True,
+            dynamo=dynamo,
+        )
+    print(f"Exported -> {output}  ({output.stat().st_size / 1e6:.1f} MB)")
+    return wrapper, output_names
+
+
+def check_parity(
+    wrapper: nn.Module, onnx_path: Path, output_names: list[str], batch: int = 2, atol: float = 1e-3
+) -> None:
+    """Compare ONNX Runtime against PyTorch on random input.
+
+    Note: this is only meaningful with the real pretrained weights. With randomly
+    initialised weights the encoder emits thousands of identical proposal scores,
+    so TopK query selection is arbitrary and torch/ORT legitimately disagree.
+    """
+    import onnxruntime as ort
+
+    x = torch.randn(batch, 3, INPUT_SIZE, INPUT_SIZE)
+    with torch.inference_mode():
+        torch_out = [t.numpy() for t in wrapper(x)]
+
+    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
+    onnx_out = sess.run(None, {"pixel_values": x.numpy()})
+
+    ok = True
+    for name, t, o in zip(output_names, torch_out, onnx_out):
+        diff = np.abs(t - o).max()
+        status = "OK " if diff < atol else "FAIL"
+        ok &= diff < atol
+        print(f"  {status} {name:14s} shape={tuple(o.shape)} max|diff|={diff:.3e}")
+    print("Parity check passed\n" if ok else "Parity check FAILED — do not ship this graph\n")
+
+
+def to_fp16(onnx_path: Path) -> Path:
+    """Half-precision copy. Keep normalisation/mask ops in fp32 to avoid overflow."""
+    import onnx
+    from onnxconverter_common import float16
+
+    model = onnx.load(str(onnx_path))
+    fp16_model = float16.convert_float_to_float16(
+        model, keep_io_types=True, op_block_list=["GridSample", "ReduceMean", "Pow", "Sqrt", "Div"]
+    )
+    out = onnx_path.with_name(onnx_path.stem + "_fp16.onnx")
+    onnx.save(fp16_model, str(out))
+    print(f"fp16 -> {out}  ({out.stat().st_size / 1e6:.1f} MB)")
+    return out
+
+
+def main() -> int:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", default=DEFAULT_MODEL)
+    p.add_argument("--output", type=Path, default=Path("pp_doclayoutv3.onnx"))
+    p.add_argument("--opset", type=int, default=17, help=">=16 required for GridSample")
+    p.add_argument("--no-masks", dest="with_masks", action="store_false")
+    p.add_argument("--static-batch", dest="dynamic_batch", action="store_false")
+    p.add_argument("--dynamo", action="store_true", help="use the TorchDynamo exporter")
+    p.add_argument("--fp16", action="store_true", help="also emit an fp16 copy")
+    p.add_argument("--skip-check", action="store_true")
+    args = p.parse_args()
+
+    wrapper, names = export(
+        args.model,
+        args.output,
+        with_masks=args.with_masks,
+        opset=args.opset,
+        dynamo=args.dynamo,
+        dynamic_batch=args.dynamic_batch,
+    )
+    if not args.skip_check:
+        check_parity(wrapper, args.output, names, batch=2 if args.dynamic_batch else 1)
+    if args.fp16:
+        to_fp16(args.output)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())