"""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())