README.md

---
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.

On the measured RTX 5090, one-page CUDA inference reserved 1048 MiB of GPU
memory for this process and peaked at 2471 MiB of resident host RAM. The CUDA
graph is 127.1 MiB. The Linux ONNX extra installs 2.94 GiB of runtime libraries.
Measurement conditions and larger batches are in [Memory](#memory).

## 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.

## Memory

Observed on one RTX 5090 (32607 MiB, driver 615.71.09) with Python 3.12.13,
ONNX Runtime GPU 1.23.2, CUDA runtime 12.8.90 and cuDNN 9.10.2.21. Stored
weights remain FP32. The 36.7M parameters occupy about 140 MiB; the GPU figures
below are larger because they include activations and the CUDA workspace.
Another GPU, driver or concurrent process can reserve a different amount.
Smaller GPUs were not tested. These rows are not a latency or accuracy result.

### One-page API

`backend="onnx"` keeps one CUDA session after a startup probe. The input was a
synthetic blank page. Every page is letterboxed to 1024 x 1024, so this GPU
reservation does not grow with the source raster's pixel count. `nvidia-smi`
reported 10 MiB in use before the process started. A second process reported
the same 784 MiB and 1048 MiB process reservations.

| Stage | Process GPU memory | `nvidia-smi` memory.used |
|---|---:|---:|
| Session loaded | 784 MiB | 802 MiB |
| After public `predict` calls | 1048 MiB | 1066 MiB |

That process peaked at 2471 MiB resident host RAM (2,530,632 KiB) and was still
at 2245 MiB after inference. CUDA maps a much larger virtual address range.
Virtual size is not RAM that has to be free.

### Raw ONNX batch

The graph input `images` has a dynamic batch `N`. The public Python API still
submits one page. Two fresh sessions, with the published CUDA TF32 options and
synthetic 1024 inputs, reserved:

| Batch | Higher process GPU memory | Higher `nvidia-smi` memory.used |
|---|---:|---:|
| 1 | 1054 MiB | 1072 MiB |
| 2 | 2582 MiB | 2600 MiB |
| 4 | 2590 MiB | 2608 MiB |
| 8 | 3614 MiB | 3632 MiB |

Batch 2 reserved 1566 MiB in one session and 2582 MiB in the other. The CUDA
workspace choice is not stable across sessions, so plan for the higher figure.
Batch 4 stayed within 8 MiB of that higher batch-2 reservation. These samples
do not qualify accuracy or throughput above batch 1.

### Disk

| Item | Size |
|---|---:|
| `model.onnx`, the CUDA graph | 127.1 MiB (133,276,847 bytes) |
| `model.safetensors`, CPU PyTorch comparison only | 140.0 MiB (146,766,504 bytes) |
| Whole repository | 267.1 MiB (280,043,351 bytes) |
| Installed Linux x86-64 `inference[onnx]` runtime libraries | 2.94 GiB (3,154,851,060 bytes) |

The CUDA path does not read `model.safetensors`. The installed library figure
is ONNX Runtime GPU 1.23.2 plus the pinned CUDA 12.8 and cuDNN 9.10 packages
on this machine. Pip printed about 1.99 GB for those wheel files before they
were unpacked.

## 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 averages IoU thresholds 0.50:0.95; AP50, AP70 and AP90 use their
respective IoU thresholds. Scores are percentages. Evaluation definitions
are private, so these numbers are not independently reproducible benchmarks.
The public prediction API uses a 0.30 score floor.

Latency measures batch-one decoded RGB to CPU box rows, including
preprocessing, transfers, inference, validation and postprocessing. It excludes
image decoding, IPC, initialization, evaluation and API serialization.
Application latency varies with host load. The measured median reduction is
44.2% versus the eager reference. Small AP differences do not establish an
accuracy improvement.

Two additional internal evaluations showed AP changes of -0.0639 and +0.0730
points. The largest class AP decrease in the primary evaluation was 0.0328
points. Validate on representative inputs before deployment.

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.

## Published inference runtime

The CUDA-only inference package does not require PyTorch. Its internal
development accuracy measurements are:

| AP | AP50 | AP70 | AP90 |
|---:|---:|---:|---:|
| 76.2396 | 90.9855 | 86.8767 | 55.9099 |

These results are not independent acceptance. Independent CUDA invocations
are not bitwise equivalent.

## 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. Observed memory for the one-page API and for raw batches 1, 2, 4 and 8 is reported under Memory.

## 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

## Quick start — CUDA + TF32

Requires supported Linux x86-64/NVIDIA CUDA setup and Python 3.12 or 3.13.
Install the download helper, then the immutable inference source:

```bash
python -m pip install "huggingface-hub>=0.36,<2"
```

```python
import subprocess
import sys
from huggingface_hub import snapshot_download

source = snapshot_download(
    "AutomatosX/AXDocLayout-M-v1.2-CUDA-TF32",
    revision="7fb4118b83a07f3af901ee9d52f3f3af67224c35",
    allow_patterns=["inference/**"],
    token=False,
)
subprocess.check_call([
    sys.executable, "-m", "pip", "install", source + "/inference[onnx]"
])
```

Run CUDA-powered layout detection:

```python
import numpy as np
from PIL import Image
from ax_doclayout import AXDocLayout, PageImage

model = AXDocLayout.from_pretrained(
    "AutomatosX/AXDocLayout-M-v1.2-CUDA-TF32",
    revision="7d2e54941a3c2c9d1ac4af27f21531bf39dfb1dd",
    backend="onnx",  # v1.2: CUDAExecutionProvider with TF32 enabled
)
page = PageImage(
    rgb=np.asarray(Image.open("page.png").convert("RGB"), dtype=np.uint8),
    page_index=0,
)
result = model.predict(page)
print(result.to_json())
```

For the separate CPU PyTorch comparison, install `inference[torch]` and
explicitly select `backend="pytorch"`. The CUDA ONNX loader does not switch
to this backend automatically.