From ebc54830ec213db0aae9b83cf422d90172c89436 Mon Sep 17 00:00:00 2001 From: wang-xinyu Date: Tue, 20 Dec 2022 13:42:59 +0800 Subject: [PATCH] refactor yolov5 dir --- README.md | 2 +- yolov5/CMakeLists.txt | 58 +++++++++---------- yolov5/README.md | 24 ++++---- yolov5/gen_wts.py | 5 +- yolov5/{ => plugin}/yololayer.cu | 0 yolov5/{ => plugin}/yololayer.h | 0 yolov5/{ => src}/calibrator.cpp | 0 yolov5/{ => src}/calibrator.h | 0 yolov5/{ => src}/common.hpp | 0 yolov5/{ => src}/cuda_utils.h | 0 yolov5/{ => src}/logging.h | 0 yolov5/{ => src}/macros.h | 0 yolov5/{ => src}/preprocess.cu | 0 yolov5/{ => src}/preprocess.h | 0 yolov5/{ => src}/utils.h | 0 yolov5/yolov5_cls.cpp | 4 +- yolov5/yolov5_cls_trt.py | 2 +- yolov5/{yolov5.cpp => yolov5_det.cpp} | 4 +- ...da_python.py => yolov5_det_cuda_python.py} | 0 yolov5/{yolov5_trt.py => yolov5_det_trt.py} | 0 yolov5/yolov5_seg.cpp | 4 +- 21 files changed, 48 insertions(+), 55 deletions(-) rename yolov5/{ => plugin}/yololayer.cu (100%) rename yolov5/{ => plugin}/yololayer.h (100%) rename yolov5/{ => src}/calibrator.cpp (100%) rename yolov5/{ => src}/calibrator.h (100%) rename yolov5/{ => src}/common.hpp (100%) rename yolov5/{ => src}/cuda_utils.h (100%) rename yolov5/{ => src}/logging.h (100%) rename yolov5/{ => src}/macros.h (100%) rename yolov5/{ => src}/preprocess.cu (100%) rename yolov5/{ => src}/preprocess.h (100%) rename yolov5/{ => src}/utils.h (100%) rename yolov5/{yolov5.cpp => yolov5_det.cpp} (98%) rename yolov5/{yolov5_trt_cuda_python.py => yolov5_det_cuda_python.py} (100%) rename yolov5/{yolov5_trt.py => yolov5_det_trt.py} (100%) diff --git a/README.md b/README.md index 416b50a..e6c8261 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ The basic workflow of TensorRTx is: ## News -- `18 Dec 2022`. [YOLOv5](./yolov5) upgrade to support v7.0, including instance segmention. +- `18 Dec 2022`. [YOLOv5](./yolov5) upgrade to support v7.0, including instance segmentation. - `12 Dec 2022`. [East-Face](https://github.com/East-Face): [UNet](./unet) upgrade to support v3.0 of [Pytorch-UNet](https://github.com/milesial/Pytorch-UNet). - `26 Oct 2022`. [ausk](https://github.com/ausk): YoloP(You Only Look Once for Panopitic Driving Perception). - `19 Sep 2022`. [QIANXUNZDL123](https://github.com/QIANXUNZDL123) and [lindsayshuo](https://github.com/lindsayshuo): YOLOv7. diff --git a/yolov5/CMakeLists.txt b/yolov5/CMakeLists.txt index 05630f3..a1522ce 100644 --- a/yolov5/CMakeLists.txt +++ b/yolov5/CMakeLists.txt @@ -1,4 +1,4 @@ -cmake_minimum_required(VERSION 2.6) +cmake_minimum_required(VERSION 3.10) project(yolov5) @@ -8,51 +8,45 @@ option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) set(CMAKE_CXX_STANDARD 11) set(CMAKE_BUILD_TYPE Debug) -find_package(CUDA REQUIRED) - -if(WIN32) +# TODO(Call for PR): make cmake compatible with Windows +set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) enable_language(CUDA) -endif(WIN32) -include_directories(${PROJECT_SOURCE_DIR}/include) # include and link dirs of cuda and tensorrt, you need adapt them if yours are different # cuda include_directories(/usr/local/cuda/include) link_directories(/usr/local/cuda/lib64) # tensorrt # TODO(Call for PR): make TRT path configurable from command line -include_directories(/usr/include/x86_64-linux-gnu/) -link_directories(/usr/lib/x86_64-linux-gnu/) +include_directories(/home/nvidia/TensorRT-8.2.5.1/include/) +link_directories(/home/nvidia/TensorRT-8.2.5.1/lib/) -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -g -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED") -cuda_add_library(myplugins SHARED yololayer.cu) +include_directories(${PROJECT_SOURCE_DIR}/src/) +include_directories(${PROJECT_SOURCE_DIR}/plugin/) +file(GLOB_RECURSE SRCS ${PROJECT_SOURCE_DIR}/src/*.cpp ${PROJECT_SOURCE_DIR}/src/*.cu) +file(GLOB_RECURSE PLUGIN_SRCS ${PROJECT_SOURCE_DIR}/plugin/*.cu) + +add_library(myplugins SHARED ${PLUGIN_SRCS}) target_link_libraries(myplugins nvinfer cudart) find_package(OpenCV) include_directories(${OpenCV_INCLUDE_DIRS}) -cuda_add_executable(yolov5 calibrator.cpp yolov5.cpp preprocess.cu) +add_executable(yolov5_det yolov5_det.cpp ${SRCS}) +target_link_libraries(yolov5_det nvinfer) +target_link_libraries(yolov5_det cudart) +target_link_libraries(yolov5_det myplugins) +target_link_libraries(yolov5_det ${OpenCV_LIBS}) -target_link_libraries(yolov5 nvinfer) -target_link_libraries(yolov5 cudart) -target_link_libraries(yolov5 myplugins) -target_link_libraries(yolov5 ${OpenCV_LIBS}) +add_executable(yolov5_cls yolov5_cls.cpp ${SRCS}) +target_link_libraries(yolov5_cls nvinfer) +target_link_libraries(yolov5_cls cudart) +target_link_libraries(yolov5_cls myplugins) +target_link_libraries(yolov5_cls ${OpenCV_LIBS}) -add_executable(yolov5-cls calibrator.cpp yolov5_cls.cpp) - -target_link_libraries(yolov5-cls nvinfer) -target_link_libraries(yolov5-cls cudart) -target_link_libraries(yolov5-cls myplugins) -target_link_libraries(yolov5-cls ${OpenCV_LIBS}) - -cuda_add_executable(yolov5-seg calibrator.cpp yolov5_seg.cpp preprocess.cu) - -target_link_libraries(yolov5-seg nvinfer) -target_link_libraries(yolov5-seg cudart) -target_link_libraries(yolov5-seg myplugins) -target_link_libraries(yolov5-seg ${OpenCV_LIBS}) - -if(UNIX) -add_definitions(-O2 -pthread) -endif(UNIX) +add_executable(yolov5_seg yolov5_seg.cpp ${SRCS}) +target_link_libraries(yolov5_seg nvinfer) +target_link_libraries(yolov5_seg cudart) +target_link_libraries(yolov5_seg myplugins) +target_link_libraries(yolov5_seg ${OpenCV_LIBS}) diff --git a/yolov5/README.md b/yolov5/README.md index a315647..94113e4 100644 --- a/yolov5/README.md +++ b/yolov5/README.md @@ -81,14 +81,14 @@ cd build cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/build cmake .. make -sudo ./yolov5 -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file -sudo ./yolov5 -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed. +sudo ./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file +sudo ./yolov5_det -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed. // For example yolov5s -sudo ./yolov5 -s yolov5s.wts yolov5s.engine s -sudo ./yolov5 -d yolov5s.engine ../samples +sudo ./yolov5_det -s yolov5s.wts yolov5s.engine s +sudo ./yolov5_det -d yolov5s.engine ../samples // For example Custom model with depth_multiple=0.17, width_multiple=0.25 in yolov5.yaml -sudo ./yolov5 -s yolov5_custom.wts yolov5.engine c 0.17 0.25 -sudo ./yolov5 -d yolov5.engine ../samples +sudo ./yolov5_det -s yolov5_custom.wts yolov5.engine c 0.17 0.25 +sudo ./yolov5_det -d yolov5.engine ../samples ``` 3. check the images generated, as follows. _zidane.jpg and _bus.jpg @@ -98,10 +98,10 @@ sudo ./yolov5 -d yolov5.engine ../samples ``` // install python-tensorrt, pycuda, etc. // ensure the yolov5s.engine and libmyplugins.so have been built -python yolov5_trt.py +python yolov5_det_trt.py // Another version of python script, which is using CUDA Python instead of pycuda. -python yolov5_trt_cuda_python.py +python yolov5_det_trt_cuda_python.py ```

@@ -115,20 +115,20 @@ python yolov5_trt_cuda_python.py wget https://github.com/joannzhang00/ImageNet-dataset-classes-labels/blob/main/imagenet_classes.txt # Build and serialize TensorRT engine -./yolov5-cls -s yolov5s-cls.wts yolov5s-cls.engine s +./yolov5_cls -s yolov5s-cls.wts yolov5s-cls.engine s # Run inference -./yolov5-cls -d yolov5s-cls.engine ../samples +./yolov5_cls -d yolov5s-cls.engine ../samples ``` ### Instance Segmentation ``` # Build and serialize TensorRT engine -./yolov5-seg -s yolov5s-seg.wts yolov5s-seg.engine s +./yolov5_seg -s yolov5s-seg.wts yolov5s-seg.engine s # Run inference -./yolov5-seg -d yolov5s-seg.engine ../samples +./yolov5_seg -d yolov5s-seg.engine ../samples ```

diff --git a/yolov5/gen_wts.py b/yolov5/gen_wts.py index 1b1a980..789e5d7 100644 --- a/yolov5/gen_wts.py +++ b/yolov5/gen_wts.py @@ -30,11 +30,10 @@ def parse_args(): pt_file, wts_file, m_type = parse_args() print(f'Generating .wts for {m_type} model') -# Initialize -device = select_device('cpu') # Load model print(f'Loading {pt_file}') -model = torch.load(pt_file, map_location=device) # load to FP32 +device = select_device('cpu') +model = torch.load(pt_file, map_location=device) # Load FP32 weights model = model['ema' if model.get('ema') else 'model'].float() if m_type in ['detect', 'seg']: diff --git a/yolov5/yololayer.cu b/yolov5/plugin/yololayer.cu similarity index 100% rename from yolov5/yololayer.cu rename to yolov5/plugin/yololayer.cu diff --git a/yolov5/yololayer.h b/yolov5/plugin/yololayer.h similarity index 100% rename from yolov5/yololayer.h rename to yolov5/plugin/yololayer.h diff --git a/yolov5/calibrator.cpp b/yolov5/src/calibrator.cpp similarity index 100% rename from yolov5/calibrator.cpp rename to yolov5/src/calibrator.cpp diff --git a/yolov5/calibrator.h b/yolov5/src/calibrator.h similarity index 100% rename from yolov5/calibrator.h rename to yolov5/src/calibrator.h diff --git a/yolov5/common.hpp b/yolov5/src/common.hpp similarity index 100% rename from yolov5/common.hpp rename to yolov5/src/common.hpp diff --git a/yolov5/cuda_utils.h b/yolov5/src/cuda_utils.h similarity index 100% rename from yolov5/cuda_utils.h rename to yolov5/src/cuda_utils.h diff --git a/yolov5/logging.h b/yolov5/src/logging.h similarity index 100% rename from yolov5/logging.h rename to yolov5/src/logging.h diff --git a/yolov5/macros.h b/yolov5/src/macros.h similarity index 100% rename from yolov5/macros.h rename to yolov5/src/macros.h diff --git a/yolov5/preprocess.cu b/yolov5/src/preprocess.cu similarity index 100% rename from yolov5/preprocess.cu rename to yolov5/src/preprocess.cu diff --git a/yolov5/preprocess.h b/yolov5/src/preprocess.h similarity index 100% rename from yolov5/preprocess.h rename to yolov5/src/preprocess.h diff --git a/yolov5/utils.h b/yolov5/src/utils.h similarity index 100% rename from yolov5/utils.h rename to yolov5/src/utils.h diff --git a/yolov5/yolov5_cls.cpp b/yolov5/yolov5_cls.cpp index 15c8619..c4bf1b8 100644 --- a/yolov5/yolov5_cls.cpp +++ b/yolov5/yolov5_cls.cpp @@ -212,8 +212,8 @@ int main(int argc, char** argv) { std::string img_dir; if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir)) { std::cerr << "arguments not right!" << std::endl; - std::cerr << "./yolov5-cls -s [.wts] [.engine] [n/s/m/l/x or c gd gw] // serialize model to plan file" << std::endl; - std::cerr << "./yolov5-cls -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; + std::cerr << "./yolov5_cls -s [.wts] [.engine] [n/s/m/l/x or c gd gw] // serialize model to plan file" << std::endl; + std::cerr << "./yolov5_cls -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; return -1; } diff --git a/yolov5/yolov5_cls_trt.py b/yolov5/yolov5_cls_trt.py index 510c5da..c2ff932 100644 --- a/yolov5/yolov5_cls_trt.py +++ b/yolov5/yolov5_cls_trt.py @@ -216,7 +216,7 @@ class warmUpThread(threading.Thread): if __name__ == "__main__": # load custom plugin and engine - engine_file_path = "build/yolov5s_cls.engine" + engine_file_path = "build/yolov5s-cls.engine" if len(sys.argv) > 1: engine_file_path = sys.argv[1] diff --git a/yolov5/yolov5.cpp b/yolov5/yolov5_det.cpp similarity index 98% rename from yolov5/yolov5.cpp rename to yolov5/yolov5_det.cpp index 70b61a7..e9f8162 100644 --- a/yolov5/yolov5.cpp +++ b/yolov5/yolov5_det.cpp @@ -305,8 +305,8 @@ int main(int argc, char** argv) { std::string img_dir; if (!parse_args(argc, argv, wts_name, engine_name, is_p6, gd, gw, img_dir)) { std::cerr << "arguments not right!" << std::endl; - std::cerr << "./yolov5 -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file" << std::endl; - std::cerr << "./yolov5 -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; + std::cerr << "./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file" << std::endl; + std::cerr << "./yolov5_det -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; return -1; } diff --git a/yolov5/yolov5_trt_cuda_python.py b/yolov5/yolov5_det_cuda_python.py similarity index 100% rename from yolov5/yolov5_trt_cuda_python.py rename to yolov5/yolov5_det_cuda_python.py diff --git a/yolov5/yolov5_trt.py b/yolov5/yolov5_det_trt.py similarity index 100% rename from yolov5/yolov5_trt.py rename to yolov5/yolov5_det_trt.py diff --git a/yolov5/yolov5_seg.cpp b/yolov5/yolov5_seg.cpp index 8977acf..5ea3476 100644 --- a/yolov5/yolov5_seg.cpp +++ b/yolov5/yolov5_seg.cpp @@ -287,8 +287,8 @@ int main(int argc, char** argv) { std::string img_dir; if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir)) { std::cerr << "arguments not right!" << std::endl; - std::cerr << "./yolov5-seg -s [.wts] [.engine] [n/s/m/l/x or c gd gw] // serialize model to plan file" << std::endl; - std::cerr << "./yolov5-seg -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; + std::cerr << "./yolov5_seg -s [.wts] [.engine] [n/s/m/l/x or c gd gw] // serialize model to plan file" << std::endl; + std::cerr << "./yolov5_seg -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; return -1; }