Metadata-Version: 2.4
Name: r3alai
Version: 2.0.1
Summary: R3AL.AI vision quantization SDK: quantize any vision model on the R3AL platform
Project-URL: Homepage, https://github.com/R3AL-AI/SDK
Project-URL: Repository, https://github.com/R3AL-AI/SDK
Author: R3AL.ai
License-Expression: MIT
License-File: LICENSE
Keywords: onnx,ptq,qat,quantization,vision,yolo
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: numpy>=1.24.0
Requires-Dist: pydantic>=2.0
Provides-Extra: all
Requires-Dist: dill>=0.3.8; extra == 'all'
Requires-Dist: gitpython>=3.1.0; extra == 'all'
Requires-Dist: matplotlib>=3.3.0; extra == 'all'
Requires-Dist: onnx>=1.15.0; extra == 'all'
Requires-Dist: onnxruntime-gpu>=1.16.0; (sys_platform != 'darwin') and extra == 'all'
Requires-Dist: onnxruntime>=1.16.0; (sys_platform == 'darwin') and extra == 'all'
Requires-Dist: onnxscript>=0.1.0; extra == 'all'
Requires-Dist: opencv-python>=4.6.0; extra == 'all'
Requires-Dist: paddle2onnx>=1.3.0; extra == 'all'
Requires-Dist: paddlepaddle>=3.0.0; extra == 'all'
Requires-Dist: pandas>=1.1.0; extra == 'all'
Requires-Dist: pillow>=10.0.0; extra == 'all'
Requires-Dist: pyyaml>=5.3.1; extra == 'all'
Requires-Dist: scipy>=1.4.0; extra == 'all'
Requires-Dist: seaborn>=0.11.0; extra == 'all'
Requires-Dist: tensorflow>=2.12.0; extra == 'all'
Requires-Dist: tf2onnx>=1.16.0; extra == 'all'
Requires-Dist: tflite2onnx>=0.4.0; extra == 'all'
Requires-Dist: thop>=0.1.1; extra == 'all'
Requires-Dist: torch>=2.1.0; extra == 'all'
Requires-Dist: tqdm>=4.66.0; extra == 'all'
Requires-Dist: ultralytics>=8.4.0; extra == 'all'
Provides-Extra: dev
Requires-Dist: ruff>=0.1.0; extra == 'dev'
Provides-Extra: export-paddle
Requires-Dist: paddle2onnx>=1.3.0; extra == 'export-paddle'
Requires-Dist: paddlepaddle>=3.0.0; extra == 'export-paddle'
Provides-Extra: export-tf
Requires-Dist: tensorflow>=2.12.0; extra == 'export-tf'
Requires-Dist: tf2onnx>=1.16.0; extra == 'export-tf'
Provides-Extra: export-tflite
Requires-Dist: tflite2onnx>=0.4.0; extra == 'export-tflite'
Provides-Extra: export-torch
Requires-Dist: onnx>=1.15.0; extra == 'export-torch'
Requires-Dist: onnxscript>=0.1.0; extra == 'export-torch'
Requires-Dist: torch>=2.1.0; extra == 'export-torch'
Provides-Extra: vision
Requires-Dist: onnx>=1.15.0; extra == 'vision'
Requires-Dist: onnxruntime-gpu>=1.16.0; (sys_platform != 'darwin') and extra == 'vision'
Requires-Dist: onnxruntime>=1.16.0; (sys_platform == 'darwin') and extra == 'vision'
Requires-Dist: pillow>=10.0.0; extra == 'vision'
Provides-Extra: yolo
Requires-Dist: dill>=0.3.8; extra == 'yolo'
Requires-Dist: gitpython>=3.1.0; extra == 'yolo'
Requires-Dist: matplotlib>=3.3.0; extra == 'yolo'
Requires-Dist: opencv-python>=4.6.0; extra == 'yolo'
Requires-Dist: pandas>=1.1.0; extra == 'yolo'
Requires-Dist: pyyaml>=5.3.1; extra == 'yolo'
Requires-Dist: scipy>=1.4.0; extra == 'yolo'
Requires-Dist: seaborn>=0.11.0; extra == 'yolo'
Requires-Dist: thop>=0.1.1; extra == 'yolo'
Requires-Dist: tqdm>=4.66.0; extra == 'yolo'
Requires-Dist: ultralytics>=8.4.0; extra == 'yolo'
Description-Content-Type: text/markdown

# r3alai

**R3AL.AI** is the brand. **`r3alai`** is the Python package (`pip install r3alai`); import it as **`r3alai`**.

Universal vision-model quantization, architecture-agnostic. Full round trip:

```
your_model (.pt/.h5/.tflite/...)  →  export to ONNX  →  quantize (PTQ) or qat_pipeline  →  quantized .onnx
                                                                                          ↳  export_native=True → back to a native .pt (torch.nn.Module, any architecture)
```

No vendor lock-in: export from PyTorch, TensorFlow/Keras, Ultralytics YOLO, TFLite, Paddle, or a ready-made ONNX file. After quantization the result can be rebuilt into a plain `.pt` file, regardless of the source architecture (ResNet, EfficientNet, ViT, YOLO, arbitrary HuggingFace vision models, ...).

Documentation: [docs.r3al.ai](https://docs.r3al.ai)

## Install

```bash
pip install "r3alai[vision]"        # PTQ + QAT on ONNX models
pip install "r3alai[vision,yolo]"   # + Ultralytics YOLO support
pip install "r3alai[all]"           # all export adapters (TF, Paddle, TFLite)
```

For local development: `pip install -e ".[vision]"`. QAT requires `onnx2torch` (included in `[vision]`).

## Quick start

```python
from r3alai.quant import Quantizer, QuantConfig

# PTQ: any vision ONNX (int8_dynamic is the default method)
Quantizer().quantize("efficientnet_b0.onnx", output_dir="./out")

# PTQ: static INT8 with calibration images (fastest, covers Conv layers)
Quantizer(QuantConfig(method="int8_static")).quantize(
    "yolov8n.onnx",
    calibration_data=["img1.jpg", "img2.jpg"],
    output_dir="./out",
)

# QAT: ONNX + training images (generic pipeline)
Quantizer(QuantConfig(mode="qat")).train_qat(
    "yolov8n.onnx",
    calibration_data=["img1.jpg", "img2.jpg"],
    output_dir="./qat_out",
    epochs=3,
    export_native=True,
    int8_runtime=True,  # also produce a genuinely smaller int8 deliverable
)

# QAT: native Ultralytics .pt in (keeps the YOLO detect/pose head intact)
Quantizer(QuantConfig(mode="qat")).train_qat(
    "yolo11n-pose.pt",
    calibration_data=["img1.jpg", "img2.jpg"],
    output_dir="./qat_out",
    source="ultralytics",
    epochs=3,
    export_native=True,  # uses YOLO.save() -> loadable with YOLO(path)
)
```

## Two paradigms

| Paradigm | Action | Input | Training | When |
|-----------|--------|-------|----------|------|
| **PTQ** (post-training) | `quantize` | `.onnx` | No | Fast, no dataset needed (except `int8_static`) |
| **QAT** (quantization-aware training) | `qat_pipeline` | `.onnx` + images | Yes (`epochs`) | PTQ accuracy not good enough, low bit widths |

## PTQ methods (`action: quantize`)

| `method` | For | Produces a genuinely smaller int8 graph? |
|----------|-----|-------------------------------------------|
| **`int8_dynamic`** | Default. ONNX Runtime dynamic INT8 (MatMul/Gemm; Conv is deliberately skipped, no reliable ConvInteger kernel) | Yes |
| `int8_static` | ONNX Runtime static INT8 QDQ (+ calibration images), covers Conv | Yes |

`int8_static` uses `per_channel=True` by default and automatically excludes ops close to the graph outputs from quantization (e.g. the Sigmoid/Concat of a detection head). With few calibration images those sensitive tail ops can otherwise collapse accuracy entirely (mAP ≈ 0). Override with `nodes_to_exclude` or `per_channel=False` if needed.

Both `int8_static` and `int8_dynamic` also exclude Softmax/LayerNormalization/Gelu/Erf (plus the MatMul/Gemm feeding directly into or out of a Softmax) from quantization by default: `exclude_attention_sensitive_ops=True`. Attention softmax and LayerNorm activations have very peaked/small-variance value ranges that saturate under a linear int8 scale; on transformer-style backbones (CLIP ViT, CLIPSeg decoders and the like) this exclusion is the difference between a working model and a full collapse. Set `exclude_attention_sensitive_ops=False` to disable, or pass an explicit `nodes_to_exclude` to override both auto-detections.

## QAT (`action: qat_pipeline`)

Universal ONNX in → train with fake-quant layers → ONNX out. No separate `method`: use `wbit`, `abit`, `epochs`.

```python
from r3alai.quant import Quantizer, QuantConfig

config = QuantConfig(mode="qat", qat_wbit=8, qat_abit=8)
result = Quantizer(config).train_qat(
    "yolov8n.onnx",
    calibration_data=["img1.jpg"],
    output_dir="./qat_out",
    epochs=1,
)
```

## Export to ONNX (any architecture)

`export_to_onnx` picks the right adapter automatically based on file extension, or force one with `source=`:

| `source` | Input | Aliases |
|----------|-------|---------|
| `ultralytics` | YOLO `.pt` | `yolo` |
| `pytorch` | saved `torch.nn.Module` (`.pt`/`.pth`, needs `input_shape`) | `torch` |
| `tensorflow` | SavedModel dir, `.h5`, `.keras`, `.pb` | `tf`, `keras` |
| `paddle` | Paddle inference model (`.pdmodel`/`.json` + `.pdiparams`) | `paddlepaddle` |
| `tflite` | `.tflite` | |
| `onnx` | validates/stages an existing `.onnx` file | |

```python
from r3alai.quant.export import export_to_onnx

export_to_onnx("yolo11n-pose.pt", output_dir="./out", source="ultralytics", imgsz=640)
export_to_onnx("resnet18.pt", output_dir="./out", source="pytorch", input_shape=[1, 3, 224, 224])
export_to_onnx("model.h5", output_dir="./out", source="tensorflow")
```

## Back to native `.pt` (after PTQ or QAT)

Every PTQ backend and QAT pipeline supports `export_native=True`: the quantized ONNX graph is rebuilt into a plain `torch.nn.Module` via [`onnx2torch`](https://github.com/ENOT-AutoDL/onnx2torch) and saved as `.pt`. Architecture-agnostic, so it works for any model the adapters above can export, not just YOLO.

```python
result = Quantizer().quantize("resnet18.onnx", output_dir="./out", export_native=True)
# result.path contains both the .quantized.onnx and a reconstructed .pt model
```

Best-effort: if the reconstruction fails (e.g. exotic `QuantizeLinear`/`QLinear*` ops that onnx2torch doesn't know), the quantization job itself does not fail. The error is reported in the manifest (`native_export_error`) and the `.onnx` deliverable remains the primary, always-valid result.

QAT folds trained `QuantConv2d` layers back into plain fp32 `Conv2d` for export, so the `*_qat.quantized.onnx` file is float32 on disk. Pass `int8_runtime=True` to additionally run an ONNX Runtime static-int8 pass on top: that yields a separate, genuinely smaller deliverable (path reported as `int8_runtime_output_model` in the manifest / API response).

## Verify your install

```bash
python -c "from r3alai.quant import Quantizer, QuantConfig; print('OK')"
python scripts/verify_sdk.py    # from a source checkout: full export → PTQ → QAT → benchmark round trip
```

## RunPod API

The RunPod serverless handler and deploy tooling live in the separate [`API-SDK` repo](https://github.com/R3AL-AI/API-SDK), which installs `r3alai` as a dependency and calls the same `Quantizer`/`QuantConfig` via a JSON job payload:

```json
{
  "input": {
    "action": "quantize",
    "model": "/runpod-volume/models/your_model.onnx"
  }
}
```

See that repo's `docs/API.md` for the full API reference.

MIT
