Skip to content

Commit 3635f0f

Browse files
zgjjazgjja
committed
add Vision Transformer
1 parent 076a8af commit 3635f0f

File tree

14 files changed

+2369
-1
lines changed

14 files changed

+2369
-1
lines changed

‎.pre-commit-config.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ repos:
1010
types: [python]
1111
- id: check-added-large-files
1212
- repo: https://github.com/pre-commit/mirrors-clang-format
13-
rev: v14.0.6
13+
rev: v18.1.3
1414
hooks:
1515
- id: clang-format
1616
types_or: [c++, c, cuda]

‎README.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ The basic workflow of TensorRTx is:
1717

1818
## News
1919

20+
- `3 Mar 2026`. [zgjja](https://github.com/zgjja) Add Vision Transformer
2021
- `2 Feb 2026`. [fazligorkembal](https://github.com/fazligorkembal) Yolo26-Det, Yolo26-Obb, Yolo26-Cls
2122
- `15 Jan 2026`. [zgjja](https://github.com/zgjja) Refactor multiple old CV models to support TensorRT SDK through 7~10.
2223
- `8 Jan 2026`. [ydk61](https://github.com/ydk61): YOLOv13
@@ -136,6 +137,7 @@ Following models are implemented.
136137
| [shufflenet](./shufflenetv2) | ShuffleNet v2 with 0.5x output channels |
137138
| [squeezenet](./squeezenet) | SqueezeNet 1.1 model |
138139
| [vgg](./vgg) | VGG 11-layer model |
140+
| [ViT](./vit) | vision transformer, using weight and model from huggingface |
139141
| [yolov3-tiny](./yolov3-tiny) | weights and pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) |
140142
| [yolov3](./yolov3) | darknet-53, weights and pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) |
141143
| [yolov3-spp](./yolov3-spp) | darknet-53, weights and pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) |

‎vit/CMakeLists.txt‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
cmake_minimum_required(VERSION 3.17.0)
2+
3+
project(
4+
vit
5+
VERSION 0.1
6+
LANGUAGES C CXX CUDA)
7+
8+
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
9+
set(CMAKE_CUDA_ARCHITECTURES 80 86 89 90 100 120)
10+
endif()
11+
12+
set(CMAKE_CXX_STANDARD 20)
13+
set(CMAKE_CXX_STANDARD_REQUIRED ON)
14+
set(CMAKE_CUDA_STANDARD 17)
15+
set(CMAKE_CUDA_STANDARD_REQUIRED ON)
16+
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
17+
set(CMAKE_INCLUDE_CURRENT_DIR TRUE)
18+
19+
option(CUDA_USE_STATIC_CUDA_RUNTIME "Use static cudaruntime library" OFF)
20+
21+
find_package(Threads REQUIRED)
22+
find_package(CUDAToolkit REQUIRED)
23+
find_package(OpenCV REQUIRED)
24+
25+
if(NOT TARGET TensorRT::TensorRT)
26+
include(FindTensorRT.cmake)
27+
else()
28+
message("TensorRT has been found, skipping for ${PROJECT_NAME}")
29+
endif()
30+
31+
add_executable(${PROJECT_NAME} "${PROJECT_NAME}.cc" "cuda_allocator.cc"
32+
"profiler.cc")
33+
target_include_directories(${PROJECT_NAME} PRIVATE ${OpenCV_INCLUDE_DIRS})
34+
target_link_libraries(
35+
${PROJECT_NAME} PUBLIC Threads::Threads CUDA::cudart CUDA::cuda_driver
36+
TensorRT::TensorRT ${OpenCV_LIBS})
37+
38+
if(WIN32)
39+
set_target_properties(
40+
${PROJECT_NAME} PROPERTIES MSVC_RUNTIME_LIBRARY
41+
"MultiThreaded$<$<CONFIG:Debug>:Debug>")
42+
endif()

‎vit/FindTensorRT.cmake‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
cmake_minimum_required(VERSION 3.17.0)
2+
3+
function(_guess_path var_name required_files)
4+
set(_result "")
5+
6+
foreach(path_entry IN LISTS ARGN)
7+
if(NOT EXISTS "${path_entry}")
8+
message(DEBUG "skip non-existing path '${path_entry}'")
9+
continue()
10+
endif()
11+
12+
set(_ok TRUE)
13+
foreach(required_file IN LISTS required_files)
14+
if(NOT EXISTS "${path_entry}/${required_file}")
15+
set(_ok FALSE)
16+
message(DEBUG "'${path_entry}' missing '${required_file}'")
17+
break()
18+
endif()
19+
endforeach()
20+
21+
if(_ok)
22+
list(APPEND _result "${path_entry}")
23+
message(DEBUG "accept '${path_entry}'")
24+
else()
25+
message(DEBUG "reject '${path_entry}'")
26+
endif()
27+
endforeach()
28+
29+
if(_result STREQUAL "")
30+
message(
31+
FATAL_ERROR
32+
"_guess_path(${var_name}) failed: no valid path found. required_files='${required_files}' candidates='${ARGN}'"
33+
)
34+
endif()
35+
36+
set(${var_name}
37+
"${_result}"
38+
PARENT_SCOPE)
39+
endfunction()
40+
41+
# add library
42+
add_library(TensorRT IMPORTED INTERFACE)
43+
add_library(TensorRT::TensorRT ALIAS TensorRT)
44+
45+
set(TRT_VERSION
46+
CACHE
47+
STRING
48+
"TensorRT version, e.g. \"8.6.1.6\" or \"8.6.1.6+cuda12.0.1.011\", \"8.6.1.6.Windows10.x86_64.cuda-12.0\" etc"
49+
)
50+
51+
if(NOT TRT_VERSION STREQUAL "" AND NOT $ENV{TRT_VERSION} STREQUAL "")
52+
message(
53+
WARNING
54+
"TRT_VERSION defined by cmake and environment variable both, using the later one"
55+
)
56+
endif()
57+
58+
if(NOT $ENV{TRT_VERSION} STREQUAL "")
59+
set(TRT_VERSION $ENV{TRT_VERSION})
60+
endif()
61+
62+
string(REGEX MATCH "([0-9]+)" _match ${TRT_VERSION})
63+
set(TRT_MAJOR_VERSION "${_match}")
64+
unset(_match)
65+
66+
if(WIN32)
67+
set(TensorRT_DIR "C:/Program Files/TensorRT-${TRT_VERSION}")
68+
if(NOT EXISTS "${TensorRT_DIR}")
69+
message(FATAL_ERROR "TensorRT_DIR=${TensorRT_DIR} does not exist!")
70+
endif()
71+
72+
if(${TRT_MAJOR_VERSION} GREATER_EQUAL 10)
73+
set(_modules nvinfer_10 nvinfer_plugin_10 nvinfer_vc_plugin_10
74+
nvinfer_dispatch_10 nvinfer_lean_10)
75+
message(DEBUG "Using ${_modules}")
76+
else()
77+
set(_modules nvinfer nvinfer_plugin nvinfer_vc_plugin nvinfer_dispatch
78+
nvinfer_lean)
79+
endif()
80+
81+
set(TensorRT_LIBRARY_DIR "${TensorRT_DIR}/lib")
82+
set(TensorRT_INCLUDE_DIR "${TensorRT_DIR}/include")
83+
elseif(UNIX)
84+
string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _trt_arch)
85+
set(_trt_include_candidates)
86+
if(_trt_arch MATCHES "^(aarch64|arm64|arch64)$")
87+
set(_trt_include_candidates "/usr/include/aarch64-linux-gnu" "/usr/include"
88+
"/usr/local/cuda/targets/aarch64-linux/include")
89+
set(_trt_library_candidates
90+
"/usr/local/tensorrt/targets/aarch64-linux-gnu/lib"
91+
"/usr/lib/aarch64-linux-gnu" "/usr/lib/aarch64-linux-gnu/tegra"
92+
"/usr/lib")
93+
elseif(_trt_arch MATCHES "^(x86_64|amd64)$")
94+
set(_trt_include_candidates
95+
"/usr/local/tensorrt/targets/x86_64-linux-gnu/include"
96+
"/usr/include/x86_64-linux-gnu" "/usr/include")
97+
set(_trt_library_candidates
98+
"/usr/local/tensorrt/targets/x86_64-linux-gnu/lib"
99+
"/usr/lib/x86_64-linux-gnu" "/usr/lib")
100+
else()
101+
message(FATAL_ERROR "Unknown architecture")
102+
endif()
103+
104+
set(_modules nvinfer nvinfer_plugin)
105+
if(${TRT_MAJOR_VERSION} GREATER_EQUAL 8)
106+
list(APPEND _modules nvinfer_vc_plugin nvinfer_dispatch nvinfer_lean)
107+
endif()
108+
109+
_guess_path(TensorRT_LIBRARY_DIR "libnvinfer.so;libnvinfer_plugin.so"
110+
${_trt_library_candidates})
111+
message(STATUS "TensorRT libraries: ${TensorRT_LIBRARY_DIR}")
112+
_guess_path(TensorRT_INCLUDE_DIR "NvInfer.h" ${_trt_include_candidates})
113+
message(STATUS "TensorRT includes: ${TensorRT_INCLUDE_DIR}")
114+
endif()
115+
116+
foreach(lib IN LISTS _modules)
117+
find_library(
118+
TensorRT_${lib}_LIBRARY
119+
NAMES ${lib}
120+
HINTS ${TensorRT_LIBRARY_DIR})
121+
list(APPEND TensorRT_LIBRARIES ${TensorRT_${lib}_LIBRARY})
122+
endforeach()
123+
124+
target_link_libraries(TensorRT INTERFACE ${TensorRT_LIBRARIES})
125+
126+
message(STATUS "Found TensorRT libs: ${TensorRT_LIBRARIES}")
127+
128+
set_target_properties(
129+
TensorRT
130+
PROPERTIES C_STANDARD 17
131+
CXX_STANDARD 17
132+
POSITION_INDEPENDENT_CODE ON
133+
SKIP_BUILD_RPATH TRUE
134+
BUILD_WITH_INSTALL_RPATH TRUE
135+
INSTALL_RPATH "$ORIGIN"
136+
INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIR}")
137+
138+
unset(TRT_MAJOR_VERSION)
139+
unset(_modules)
140+
unset(_trt_include_candidates)
141+
unset(_trt_library_candidates)
142+
unset(_trt_arch)

0 commit comments

Comments
 (0)