wang-xinyu/tensorrtx

What it solves

TensorRTx provides a way to implement popular deep learning networks using the TensorRT network definition API instead of relying on standard parsers (like ONNX or UFF). This approach avoids the "black box" nature of automated parsers, allowing developers to have full control over the network structure and optimization.

How it works

The project follows a specific workflow to move from a trained model to a high-performance inference engine:

1. Weight Export: Weights are extracted from frameworks like PyTorch, MXNet, or TensorFlow and exported to a plain text .wts file.
2. Network Definition: The network is rebuilt from scratch using TensorRT's complex APIs to define the layers and structure.
3. Engine Building: The .wts weights are loaded into the defined network to build a TensorRT engine.
4. Inference: The resulting engine is loaded and executed for inference.

Who it’s for

Developers and researchers who need maximum flexibility to modify networks, debug middle layer results, or learn the internal structure of deep learning models while deploying them on NVIDIA hardware via TensorRT.

Highlights

• High Flexibility: Easily add, delete, or replace layers and integrate preprocessing and postprocessing directly into the network.
• Detailed Debugging: Supports incremental development to inspect intermediate layer outputs.
• Extensive Model Zoo: Includes implementations for a wide variety of models including YOLO series (v3 through v13), ResNet, MobileNet, Vision Transformers (ViT), and specialized models for face detection (RetinaFace) and text detection (DBNet).
• Customizable Workflow: Supports building individual subprojects or the entire suite via CMake.