🚀 OxidizedVision

Compile your PyTorch models to Rust for ultra-fast, memory-safe inference.

OxidizedVision is a production-grade toolkit that bridges the gap between Python-based model training and Rust-based deployment. It provides a seamless pipeline to convert, optimize, validate, benchmark, profile, and package your models — from a trained PyTorch nn.Module to a deployable Rust binary, REST API, or WebAssembly module.

---

✨ Key Features

Feature	Description
🔄 Model Conversion	PyTorch → TorchScript → ONNX with a single command
⚡ Optimization	ONNX graph simplification, constant folding, dynamic/static (calibration-based) INT8, FP16 quantization, magnitude pruning
✅ Validation	Numerical consistency checks (MAE, RMSE, Cosine Similarity) across formats
📊 Benchmarking	Latency (avg, p50, p95, p99), throughput, and memory profiling
🔬 Profiling	Parameter count, model size, per-layer breakdown
📦 Packaging	Auto-generate a deployable Rust crate (server or CLI)
🌐 Multi-Backend	tract (pure Rust), ort (ONNX Runtime, fused vision kernels), tch (LibTorch), tensorrt (NVIDIA GPU)
🧩 WASM Support	Run models in the browser via WebAssembly
📋 Model Registry	Track all converted models and their metadata locally
🎨 Rich CLI	Beautiful terminal output with progress indicators and tables
🔀 Multi-Model Server	Serve multiple models from a single Rust server instance
⏱️ Dynamic Batching	Configurable request batching for efficient inference
📝 Structured Logging	tracing (Rust) + Rich/JSON (Python) for full observability
📈 Metrics Endpoint	/metrics for monitoring request counts and server health

---

🏗️ Architecture

flowchart TB

subgraph CLI["Python Client (CLI)"]
    direction TB
    commands["convert | validate | benchmark | optimize | profile<br/>package | serve | list | info"]
    globals["Global Flags:<br/>--verbose | --json-log"]
end

CLI -->|Generates| RUST

subgraph RUST["Rust Runtimes"]
    direction TB

    TCH["runner_tch<br/>(LibTorch)"]
    TRACT["runner_tract<br/>(Pure Rust)"]
    ORT["runner_ort<br/>(ONNX Runtime, fused kernels)"]
    TRT["runner_tensorrt<br/>(GPU / TensorRT)"]

    CORE["runner_core (Shared Trait)<br/>+ tracing structured logging"]

    TCH --> CORE
    TRACT --> CORE
    ORT --> CORE
    TRT --> CORE
end

RUST -->|Deploys to| NATIVE
RUST --> REST
RUST --> WASM

NATIVE["Native Binary"]
REST["REST API Server<br/>(multi-model, batching, /metrics)"]
WASM["WASM Module"]

Loading

---

⚡ Quickstart

1. Install

# From PyPI
pip install oxidizedvision

# From source (development, run from the repo root)
pip install -e ".[dev]"

2. Create a Config

# config.yml
model:
  path: examples/example_unet/model.py
  class_name: UNet
  input_shape: [1, 3, 256, 256]

export:
  output_dir: out
  model_name: unet

validate:
  tolerance_mae: 1e-4
  tolerance_cos_sim: 0.999

benchmark:
  iters: 100
  device: cpu

# Optional: prune the PyTorch model before export (see "Pruning" below)
# optimize:
#   pruning_amount: 0.3       # zero out 30% of Conv/Linear weights
#   pruning_structured: false # true = zero whole output channels instead

3. Run the Pipeline

# Convert PyTorch → TorchScript + ONNX
oxidizedvision convert config.yml

# Validate numerical consistency
oxidizedvision validate config.yml

# Optimize the ONNX model
oxidizedvision optimize out/unet.onnx --quantize int8

# Benchmark performance
oxidizedvision benchmark out/unet.pt --runners torchscript,tract,ort

# Profile the model
oxidizedvision profile config.yml

# Package into a Rust crate
oxidizedvision package out/unet.onnx --runner tract --template server

# List registered models
oxidizedvision list

4. Debug with Structured Logging

# Verbose mode (DEBUG level)
oxidizedvision --verbose convert config.yml

# JSON log output (for CI / log aggregation)
oxidizedvision --json-log convert config.yml

---

📖 CLI Reference

Command	Description	Example
convert	Convert PyTorch → TorchScript + ONNX	oxidizedvision convert config.yml
validate	Check numerical consistency	oxidizedvision validate config.yml --num-tests 5
benchmark	Measure inference performance	oxidizedvision benchmark out/model.pt --runners torchscript,tract,ort
optimize	Optimize an ONNX model	oxidizedvision optimize out/model.onnx --quantize static_int8 --input-shape 1,3,256,256
profile	Analyze model parameters and layers	oxidizedvision profile config.yml
package	Generate deployable Rust crate	oxidizedvision package out/model.onnx --template server
serve	Start inference server	oxidizedvision serve ./binary --port 8080
list	List registered models	oxidizedvision list
info	Detailed model information	oxidizedvision info unet

Global Options

Flag	Description
--verbose / -v	Enable DEBUG-level logging
--json-log	Emit logs as JSON lines (for CI / production)

Pruning

Set optimize.pruning_amount (and optionally pruning_structured) in a config's YAML — pruning happens on the PyTorch model before export, so it's part of convert, not the separate optimize ONNX-stage command:

optimize:
  pruning_amount: 0.3        # zero out 30% of Conv2d/Conv1d/Linear weights
  pruning_structured: false  # false: unstructured (individual weights)
                              # true:  structured (whole output channels)

What this actually gets you today: a smaller, more compressible checkpoint (fewer/zeroed nonzero weights) and a documented sparsity percentage. It does not speed up inference on any of this repo's backends (ONNX Runtime, tract, LibTorch) — they all run dense kernels that still multiply through the zeros. Structured pruning zeros whole output channels, which is the prerequisite for actually shrinking those tensors and cutting FLOPs, but this step doesn't yet re-slice the model to drop the zeroed channels. Use it for compression today; it's the foundation real channel removal (and thus a genuine latency win) would build on.

---

🦀 Rust Runtimes

Shared Runner Trait

All backends implement a common Runner trait:

pub trait Runner: Send + Sync {
    fn from_config(config: &RunnerConfig) -> Result<Self> where Self: Sized;
    fn run(&self, input: &ArrayD<f32>) -> Result<ArrayD<f32>>;
    fn info(&self) -> ModelInfo;
}

Available Backends

Backend	Model Format	GPU	WASM	Dependencies
runner_tract	ONNX	❌	✅	None (pure Rust)
runner_ort	ONNX	✅ (CUDA/CoreML)	❌	ONNX Runtime
runner_tch	TorchScript	✅	❌	LibTorch
runner_tensorrt	ONNX → Engine	✅	❌	TensorRT SDK

runner_tensorrt shells out to trtexec (no stable Rust TensorRT bindings exist) and caches the built .engine file next to the ONNX model. Verified end-to-end (export → build engine → run inference) against a real TensorRT 11.2 SDK on a T4 GPU via benchmarks/modal_tensorrt_check.py — that run caught two trtexec CLI changes in TensorRT 10+ (bare --fp16 and --saveOutput were both removed in favor of strongly-typed networks and JSON-based --exportOutput), which are now handled with version-tolerant fallbacks so this works on both older and current TensorRT installs.

runner_ort runs models through ONNX Runtime's GraphOptimizationLevel::Level3 optimizer, which fuses common vision-backbone patterns (Conv+BatchNorm+Activation, MatMul+Add, LayerNorm, GELU) into single fused kernels and dispatches to hardware-tuned (oneDNN / cuDNN / Core ML) implementations. For standard CNN and ViT vision models this generally beats tract on both CPU and GPU latency — use tract when you need pure-Rust/WASM portability, and ort when you want the fastest native CPU/GPU inference. It holds a small pool of independent ONNX Runtime sessions (default: up to 4, sized to available parallelism) rather than one session behind a single lock, so concurrent requests in a server don't serialize onto one &mut Session.

GPU/accelerator dispatch is behind Cargo features, since each links an external SDK: cargo build -p runner_ort --features cuda for NVIDIA GPUs, --features coreml for Apple Silicon (dispatches eligible ops to the Neural Engine via Core ML, falling back to CPU for the rest) — image_server forwards the same features (cargo run -p image_server --features coreml -- ...).

---

🖥️ Inference Server

The built-in image_server example provides a production-ready REST API:

# Single model
cargo run -p image_server -- --model model.onnx --port 8080

# Multi-model (serve multiple models simultaneously)
cargo run -p image_server -- \
  --model segmenter=models/seg.onnx \
  --model classifier=models/cls.onnx \
  --port 8080

# With dynamic batching
cargo run -p image_server -- \
  --model model.onnx \
  --max-batch-size 8 \
  --max-wait-ms 50

# JSON structured logs
cargo run -p image_server -- --model model.onnx --log-format json

Endpoints

Method	Path	Description
POST	/predict	Inference on the default model from a raw flattened-tensor JSON body
POST	/predict/{model_name}	Same, on a named model
POST	/predict/image	Inference on the default model from a raw image upload (JPEG/PNG/etc) — decoded, SIMD-resized, and normalized in Rust before inference
POST	/predict/image/{model_name}	Same, on a named model
GET	/health	Health check with per-model status
GET	/metrics	Prometheus scrape endpoint (text exposition format)
GET	/metrics.json	Same numbers as /metrics, as JSON
GET	/models	List all loaded models

/predict/image* accepts the raw image bytes as the request body (e.g. curl -X POST --data-binary @photo.jpg http://localhost:8080/predict/image) and resizes to the model's configured input [H, W] using a SIMD (SSE4.1/AVX2/NEON) Lanczos3 resize kernel via fast_image_resize, fusing the u8→normalized-f32 NCHW conversion into the same pass — so preprocessing never round-trips through Python or a scalar per-pixel loop.

---

🗂️ Project Structure

Oxidized-Vision/
├── python_client/             # Python CLI & pipeline
│   ├── oxidizedvision/
│   │   ├── cli.py             # Typer CLI entry point
│   │   ├── config.py          # Pydantic config models
│   │   ├── convert.py         # Model conversion
│   │   ├── validate.py        # Numerical validation
│   │   ├── benchmark.py       # Performance measurement
│   │   ├── optimize.py        # ONNX optimization
│   │   ├── profile.py         # Model profiling
│   │   ├── registry.py        # Model registry
│   │   └── logging.py         # Structured logging (Rich / JSON)
│   └── tests/                 # pytest test suite
├── rust_runtime/              # Rust inference runtimes
│   ├── crates/
│   │   ├── runner_core/       # Shared Runner trait + tracing
│   │   ├── runner_tch/        # LibTorch backend
│   │   ├── runner_tract/      # tract (ONNX) backend
│   │   ├── runner_ort/        # ONNX Runtime backend (fused vision kernels)
│   │   └── runner_tensorrt/   # TensorRT backend
│   └── examples/
│       ├── image_server/      # Multi-model REST API with batching
│       ├── denoiser_cli/      # Image denoising CLI
│       └── wasm_frontend/     # Browser inference demo
├── tools/                     # Standalone scripts
├── benchmarks/                # Benchmark infrastructure
├── examples/                  # User-facing examples
│   ├── example_unet/          # 2D segmentation: UNet with skip connections
│   ├── example_detector/      # 2D detection: compact YOLO-style detector
│   └── example_pointnet/      # 3D vision: PointNet point-cloud classifier
├── docs/                      # Architecture docs
└── .github/workflows/         # CI/CD + PyPI auto-deploy

---

📊 Benchmark Leaderboard

Real, reproducible numbers (including a GPU run on Modal) live in benchmarks/RESULTS.md — not hand-edited, generated by the scripts in benchmarks/.

---

🧪 Testing

# Python tests
pytest python_client/tests/ -v --cov=oxidizedvision

# Rust tests
cd rust_runtime && cargo test --workspace

Pre-commit Hooks

pip install pre-commit
pre-commit install
pre-commit run --all-files

---

📄 License

MIT License — see LICENSE for details.