https://github.com/dnth/supercharge-your-pytorch-image-models-blogpost
Supercharge Your PyTorch Image Models: Bag of Tricks to 8x Faster Inference with ONNX Runtime & Optimizations
https://github.com/dnth/supercharge-your-pytorch-image-models-blogpost
computer-vision onnx onnxruntime onnxruntime-gpu pytorch tensorrt timm
Last synced: 11 months ago
JSON representation
Supercharge Your PyTorch Image Models: Bag of Tricks to 8x Faster Inference with ONNX Runtime & Optimizations
- Host: GitHub
- URL: https://github.com/dnth/supercharge-your-pytorch-image-models-blogpost
- Owner: dnth
- Created: 2024-09-11T11:12:17.000Z (almost 2 years ago)
- Default Branch: main
- Last Pushed: 2024-10-04T16:51:37.000Z (almost 2 years ago)
- Last Synced: 2024-12-02T07:45:55.901Z (over 1 year ago)
- Topics: computer-vision, onnx, onnxruntime, onnxruntime-gpu, pytorch, tensorrt, timm
- Language: Jupyter Notebook
- Homepage: https://dicksonneoh.com/portfolio/supercharge_your_pytorch_image_models/
- Size: 62 MB
- Stars: 20
- Watchers: 1
- Forks: 0
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
README

This repository contains code to optimize PyTorch image models using ONNX Runtime and TensorRT, achieving up to 8x faster inference speeds. Read the full blog post [here](https://dicksonneoh.com/portfolio/supercharge_your_pytorch_image_models/).
## Installation
Create and activate a conda environment:
```bash
conda create -n supercharge_timm_tensorrt python=3.11
conda activate supercharge_timm_tensorrt
```
Install required packages:
```bash
pip install timm
pip install onnx
pip install onnxruntime-gpu==1.19.2
pip install cupy-cuda12x
pip install tensorrt==10.1.0 tensorrt-cu12==10.1.0 tensorrt-cu12-bindings==10.1.0 tensorrt-cu12-libs==10.1.0
```
Install CUDA dependencies:
```bash
conda install -c nvidia cuda=12.2.2 cuda-tools=12.2.2 cuda-toolkit=12.2.2 cuda-version=12.2 cuda-command-line-tools=12.2.2 cuda-compiler=12.2.2 cuda-runtime=12.2.2
```
Install cuDNN:
```bash
conda install cudnn==9.2.1.18
```
Set up library paths:
```bash
export LD_LIBRARY_PATH="/home/dnth/mambaforge-pypy3/envs/supercharge_timm_tensorrt/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="/home/dnth/mambaforge-pypy3/envs/supercharge_timm_tensorrt/lib/python3.11/site-packages/tensorrt_libs:$LD_LIBRARY_PATH"
```
## Running the code
The following codes correspond to the steps in the blog post.
### Load timm model and run inference:
```bash
python 00_load_and_infer.py
```
Read more [here](https://dicksonneoh.com/portfolio/supercharge_your_pytorch_image_models/#-load-and-infer)
### PyTorch latency benchmark:
```bash
python 01_pytorch_latency_benchmark.py
```
Read more [here](https://dicksonneoh.com/portfolio/supercharge_your_pytorch_image_models//#-baseline-latency)
### Convert model to ONNX:
```bash
python 02_convert_to_onnx.py
```
Read more [here](https://dicksonneoh.com/portfolio/supercharge_your_pytorch_image_models//#-convert-to-onnx)
### ONNX Runtime CPU inference:
```bash
python 03_onnx_cpu_inference.py
```
Read more [here](https://dicksonneoh.com/portfolio/supercharge_your_pytorch_image_models//#-onnx-runtime-on-cpu)
### ONNX Runtime CUDA inference:
```bash
python 04_onnx_cuda_inference.py
```
Read more [here](https://dicksonneoh.com/portfolio/supercharge_your_pytorch_image_models//#-onnx-runtime-on-cuda)
### ONNX Runtime TensorRT inference:
```bash
python 05_onnx_trt_inference.py
```
Read more [here](https://dicksonneoh.com/portfolio/supercharge_your_pytorch_image_models//#-onnx-runtime-on-tensorrt)
### Export preprocessing to ONNX:
```bash
python 06_export_preprocessing_onnx.py
```
Read more [here](https://dicksonneoh.com/portfolio/supercharge_your_pytorch_image_models//#-bake-pre-processing-into-onnx)
### Merge preprocessing and model ONNX:
```bash
python 07_onnx_compose_merge.py
```
Read more [here](https://dicksonneoh.com/portfolio/supercharge_your_pytorch_image_models//#-bake-pre-processing-into-onnx)
### Run inference on merged model:
```bash
python 08_inference_merged_model.py
```
Read more [here](https://dicksonneoh.com/portfolio/supercharge_your_pytorch_image_models//#-bake-pre-processing-into-onnx)
### Run inference on video:
```bash
python 09_video_inference.py sample.mp4 output.mp4 --live
```
https://github.com/user-attachments/assets/1a25dd6e-3512-475a-9541-29e836022bb5