license: agpl-3.0
library_name: ax-doclayout
pipeline_tag: object-detection
tags:
- document-layout
- object-detection
- onnx
- cuda
- tf32
AXDocLayout-M v1.2 — ONNX CUDA + TF32
CUDA-powered document layout detection. AX DocLayout locates layout regions in a decoded RGB page image and returns class labels, confidence scores and bounding boxes in the original raster's pixel coordinates.
What changed in v1.2
This is a CUDA runtime release using the same FP32 weights and ONNX graph as v1.1. The ONNX backend now explicitly enables ONNX Runtime's CUDA execution provider and TF32. Postprocessing is vectorized with equivalent raster geometry. The source package verifies the CUDA provider, effective TF32 policy and placement of major compute operations. CUDA setup failures raise an error; the loader does not retry with a CPU-only session.
The architecture remains RT-DETRv2 rtdetrv2_r50vd_m, with 11 layout classes.
There is no integer quantization or model retraining in this release. Both
model.onnx (standard operators, opset 16) and FP32 model.safetensors are
provided. Complete inference source is in inference/; attribution is in NOTICE.
CUDA requirements
Use Linux x86-64, Python 3.12 or 3.13, a supported NVIDIA GPU with TF32 capability
(Ampere or newer), and a driver compatible with CUDA 12.8. The ONNX extra pins
ONNX Runtime GPU 1.23.2 and includes CUDA/cuDNN Python dependencies. The tested
configuration is RTX 5090, CUDA 12.8 and cuDNN 9.10.2.21. The ONNX path requires
CUDA. backend="pytorch" remains an explicit CPU FP32 comparison path and does
not provide the CUDA performance below.
Floating-point CUDA execution can produce small output differences between invocations or sessions. Outputs are not promised to be bitwise reproducible.
TF32 is an execution-provider setting. ONNX inputs and stored model weights remain FP32; downloading the ONNX file alone does not enable TF32.
Reference development results
The following aggregate results are from internal development validation, not an independent held-out acceptance test. These are measured results for one hardware/runtime configuration, not a latency guarantee.
| Runtime configuration | p50 ms/page | p95 ms/page | Serial pages/s | AP | AP50 | AP70 | AP90 |
|---|---|---|---|---|---|---|---|
| Prior eager PyTorch FP32 reference | 46.242 | 62.540 | 19.48 | 76.2382 | 90.9897 | 86.8844 | 55.8986 |
| Optimized ONNX CUDA FP32 | 34.423 | 40.770 | 27.76 | 76.2384 | 90.9896 | 86.8844 | 55.8988 |
| Optimized ONNX CUDA + TF32 | 25.792 | 32.395 | 35.91 | 76.2414 | 90.9891 | 86.8792 | 55.9025 |
AP uses all 11 classes on 6,489 development pages, with a 0.001 score floor for evaluation. The public prediction API uses a 0.30 score floor. Latency uses batch one, three counterbalanced passes of 100 pages after untimed startup samples. The timed region is decoded RGB to CPU box rows, including letterboxing, normalization, transfers, inference, validation and postprocessing; it excludes image-file decoding, IPC, initialization and evaluation. API serialization costs and host-load variation can change application latency. The measured p50 reduction is 44.2% versus the eager reference. Small AP differences do not establish an accuracy improvement.
On two additional financial-document development cohorts (256 and 100 pages), table AP changed by -0.0639 and +0.0730 points. The largest class AP decrease on the general development cohort was 0.0328 points. Corpus and per-page data are not included.
The tested graph placed convolutions, matrix multiplications, GridSample and Softmax on CUDA. Some shape/control and arithmetic nodes execute on CPU; this is not an all-GPU graph claim.
Release installation validation
A separate clean installation of the CUDA-only inference package, with no PyTorch installed, was evaluated on all 6,489 general development pages:
| AP | AP50 | AP70 | AP90 |
|---|---|---|---|
| 76.2396 | 90.9855 | 86.8767 | 55.9099 |
The same installation passed the predefined development quality margins on both financial-document cohorts. An additional 128-page API check verified preprocessing and exact postprocessing of the captured inference outputs. Independent CUDA invocations are not bitwise-equivalent; these checks do not claim cross-session raw-detection equivalence or independent acceptance.
Input and output
Supply one HWC uint8 RGB page raster. Preprocessing preserves aspect ratio,
resizes the long edge to 1024, pads to 1024 x 1024 with white and normalizes to
NCHW float32 in [0, 1], under rgb-letterbox-1024/1. The Python API applies
score floor 0.30, removes reserved/invalid labels, reverses letterboxing and
clips boxes to the original raster. Index 0 is reserved.
The JSON schema is ax-doclayout.detection/1, taxonomy ax-doclayout-11/1.
Boxes use half-open pixel edges with top-left origin, x rightward and y
downward. Raw ONNX inputs are images: float32 [N, 3, 1024, 1024] and
orig_target_sizes: int64 [N, 2], [1024, 1024] per image. Raw outputs are
labels, boxes (letterbox-space xyxy) and scores.
| ID | Label |
|---|---|
| 1 | caption |
| 2 | footnote |
| 3 | formula |
| 4 | list_item |
| 5 | page_footer |
| 6 | page_header |
| 7 | picture |
| 8 | section_header |
| 9 | table |
| 10 | text |
| 11 | title |
Use and limitations
Validate the pinned release on representative pages before deployment. Boundary precision, very small regions and unusual document designs can need review. Confidence scores do not guarantee correctness. The model accepts page rasters; it does not open PDFs, perform OCR, extract text, predict reading order/hierarchy, or reconstruct table cells. Larger-batch throughput and other GPU/software configurations were not qualified by these results.
License
The weights and modified inference source are AGPL-3.0-only. Commercial use
is permitted subject to the license terms. See LICENSE and NOTICE for terms,
source-availability requirements and preserved Apache-2.0 component notices.
This does not relicense previous releases. Complete inference source is supplied
in inference/; datasets and the training pipeline are not distributed.
Installation and inference
Use the immutable source and weight revisions in the root README quick-start.
Install the ONNX GPU extra and explicitly select backend="onnx" for CUDA TF32.
An explicit CPU comparison is available through inference[torch] and
backend="pytorch". The package verifies payload digests before inference.
Runtime reference: https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html