PP-DocLayoutV3-ONNX / export_pp_doclayout_v3.py

Upload ONNX export

a15d07c verified 2 months ago

8.51 kB

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

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