custom-dl-optimizer 1.1.0

A profile-guided PyTorch inference optimizer using FX and Triton

Custom-DL-Optimizer

Custom-DL-Optimizer is a lightweight, research-oriented PyTorch inference optimizer for NVIDIA GPUs. It applies profiling, PyTorch FX graph rewriting, custom Triton operator replacement, channels-last memory layout conversion, and automatic mixed precision to improve inference latency while preserving the normal PyTorch module interface.

The package is designed for ML systems experiments, compiler-oriented coursework, and prototype deployment studies where the model should remain a PyTorch nn.Module rather than being exported to an external runtime.

Highlights

• PyTorch-native API: optimize an existing torch.nn.Module in one call.
• FX graph surgery: traces model graphs and rewrites supported operators.
• Triton kernel injection: replaces supported activations with custom GPU kernels when safe.
• Hardware-aware execution: supports channels-last memory layout and FP16 autocast.
• Safe fallbacks: CPU tensors and unsupported layouts fall back to PyTorch operators.
• Research notebook: includes a Colab benchmark workflow for reproducible tables and figures.

Installation

pip install custom-dl-optimizer

For GPU acceleration, use a CUDA-enabled PyTorch build and a Linux/Colab environment with NVIDIA GPU support.

python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

Quick Start

import torch
from torchvision.models import resnet50
from custom_dl_optimizer import AutoOptimizer

device = "cuda" if torch.cuda.is_available() else "cpu"
model = resnet50(weights=None).to(device).eval()
dummy_input = torch.randn(32, 3, 224, 224, device=device)

optimizer = AutoOptimizer(model, device=device)
optimized_model = optimizer.optimize(dummy_input)

with torch.inference_mode():
    baseline_output = model(dummy_input)
    optimized_output = optimized_model(
        dummy_input.to(memory_format=torch.channels_last)
        if dummy_input.is_cuda
        else dummy_input
    )

print(torch.allclose(baseline_output.float(), optimized_output.float(), rtol=8e-2, atol=8e-2))

API

AutoOptimizer(model, device="cuda")

Creates an optimizer wrapper around a PyTorch model.

Parameters:

• model: a torch.nn.Module.
• device: target device. Use "cuda" for GPU optimization or "cpu" for fallback behavior.

optimize(dummy_input)

Runs the optimization pipeline and returns an optimized nn.Module.

Pipeline stages:

1. Profile the model with torch.profiler.
2. Trace the model into a PyTorch FX graph.
3. Replace supported modules, currently nn.ReLU, with optimized alternatives.
4. Convert CUDA models to channels-last memory format for convolution-heavy workloads.
5. Wrap CUDA inference in FP16 automatic mixed precision.

Architecture

Component	File	Responsibility
Orchestrator	custom_dl_optimizer/core/engine.py	User-facing optimization pipeline
Profiler	custom_dl_optimizer/core/profiler.py	Runtime operator profiling and bottleneck reporting
Graph Surgeon	custom_dl_optimizer/core/graph_surgeon.py	FX tracing and module replacement
Triton Kernels	custom_dl_optimizer/core/triton_kernels.py	Custom GPU kernels and PyTorch fallbacks

Benchmark Snapshot

The repository includes Custom_DL_Optimizer_Research_Colab.ipynb, a Colab benchmark notebook that generates CSV, LaTeX tables, and publication-style graphs.

Example NVIDIA Tesla T4 results from the research notebook:

Model	Batch	Eager FP32 (ms)	AMP/NHWC (ms)	Custom-DL (ms)	Speedup vs Eager	Speedup vs AMP
ResNet-50	128	430.317	161.772	109.243	3.94x	1.48x
MobileNet-V2	128	124.204	62.109	35.989	3.45x	1.73x
VGG-16	128	751.234	315.222	255.445	2.94x	1.23x
EfficientNet-B0	128	171.223	80.883	45.675	3.75x	1.77x
DenseNet-121	128	414.769	227.366	132.560	3.13x	1.72x

These results should be treated as hardware- and software-version-dependent measurements. For papers or reports, rerun the notebook on the target GPU and report the generated CSV values.

Reproducible Benchmarking

1. Open Custom_DL_Optimizer_Research_Colab.ipynb in Google Colab.
2. Select Runtime > Change runtime type > GPU.
3. Restart the runtime.
4. Run all cells.
5. Use the generated files from:

/content/custom_dl_optimizer_research_outputs

The notebook exports:

• custom_dl_optimizer_research_results.csv
• custom_dl_optimizer_research_results.json
• paper_table_results.tex
• figure_1_latency_comparison.png
• figure_2_speedup.png
• figure_3_pass_coverage.png
• figure_4_output_parity.png
• generated_abstract_results.txt

Scope and Limitations

Custom-DL-Optimizer is a research prototype, not a replacement for TensorRT, Torch-TensorRT, TVM, XLA, or TorchInductor. It is most useful as a PyTorch-native experimentation layer for studying compiler passes and GPU-aware inference transformations.

Current limitations:

• The package-level optimizer currently focuses on CNN-style graphs and nn.ReLU replacement.
• Performance depends on GPU architecture, PyTorch version, CUDA version, tensor shapes, and batch size.
• User-defined Triton kernels may interact differently with torch.compile across PyTorch/Triton versions.
• Output parity should be checked with FP16-appropriate tolerances when AMP is enabled.

Roadmap

• Add explicit TorchInductor and CUDA Graph integration to the package API.
• Add Conv-BatchNorm folding as a first-class graph pass.
• Add transformer-oriented LayerNorm/GELU and MLP block optimizations.
• Add post-training INT8 quantization experiments.
• Add CI tests for graph rewriting and output parity.
• Publish benchmark reports for T4, L4, A10, and A100 GPUs.

Development

git clone https://github.com/Devrajsinh-Jhala/Custom-DL-Optimizer.git
cd Custom-DL-Optimizer
python -m pip install -e ".[dev]"
pytest

Build the package locally:

python -m build
python -m twine check dist/*

Upload to PyPI after checking the package description:

python -m twine upload dist/*

Citation

If you use this project in a paper, report the exact package version, GPU, CUDA version, PyTorch version, batch size, warmup count, benchmark iterations, and output-parity tolerance.

Suggested citation text:

Custom-DL-Optimizer: A profile-guided PyTorch micro-compiler for hardware-aware NVIDIA GPU inference optimization.

License

This project is released under the MIT License. See LICENSE.