BiRefNet TensorRT Inference
August 30, 2024 ยท View on GitHub
Introduction
This project provides code for performing inference with BiRefNet using TensorRT. The aim is to accelerate the inference process by leveraging the high-performance capabilities of TensorRT.
Inference Time Comparison
1. First Inference Time
| Method | Pytorch | ONNX | Tensorrt |
|---|---|---|---|
| inference time | 0.71s | 5.32s | 0.17s |
2. Average Inference Time (excluding the first)
| Method | Pytorch | ONNX | Tensorrt |
|---|---|---|---|
| inference time | 0.15s | 4.43s | 0.11s |
Note:
- Both the PyTorch and ONNX models are from the official BiRefNet GitHub.
- The TensorRT model was converted using Convert-ONNX-Model-to-TensorRT-Engine.
- All tests were conducted on a Win10 system with an RTX 4080 Super.
- Refer to model_compare.py for the conversion code.
Features
- Efficient inference with BiRefNet using TensorRT
- foreground estimate
- colab example
- Performance comparison between PyTorch, ONNX, and TensorRT inference
- Inference using Docker for an isolated and reproducible environment
Prerequisites
- NVIDIA GPU with CUDA(>=11.X) and Cudnn(>=8.X)
- Python 3.9
Installation
pip install -r requirements.txt
Usage
1. download onnx model
First, download onnx model from Google Drive
2. Convert ONNX Model to TensorRT Engine
second, convert your ONNX model to a TensorRT engine using the provided conversion script:
from utils import convert_onnx_to_engine
onnx_file_path = "birefnet.onnx"
engine_file_path = "engine.trt"
convert_onnx_to_engine(onnx_file_path, engine_file_path)
3. Run Inference
Now, you can run inference using the TensorRT engine with the following command:
3.1 single infer
python .\infer.py --image-path image_path --output-path result.png --output-alpha-path result_alpha.png --engine-path .\engine.trt
3.2 infer for directory
python .\infer.py --image-path image_dir --output-path output_dir --output-alpha-path alpha_dir --engine-path .\engine.trt --mode m
Contributing
Contributions are welcome! Please feel free to submit a Pull Request or open an Issue if you have any suggestions or find bugs.