haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix and optimize

Browse Source
This commit is contained in:
wang-xinyu 2021-01-29 16:20:48 +08:00
parent 6161f28854
commit 1a82822474
9 changed files with 51 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
README.md
View File

@ -10,6 +10,7 @@ All the models are implemented in pytorch/mxnet/tensorflown first, and export a
## News
- `29 Jan 2021`. U-Net added by [YuzhouPeng](https://github.com/YuzhouPeng).
- `24 Jan 2021`. IBN-Net added by [TCHeish](https://github.com/TCHeish), PSENet optimized, YOLOv5 v4.0 INT8, etc.
- `8 Jan 2021`. YOLOv5 s/m/l/x updated to v4.0.
- `27 Dec 2020`. HRNet-Semantic-Segmentation added by [BaofengZan](https://github.com/BaofengZan).
@ -24,7 +25,6 @@ All the models are implemented in pytorch/mxnet/tensorflown first, and export a
- `28 Aug 2020`. [BaofengZan](https://github.com/BaofengZan) added a tutorial for compiling and running tensorrtx on windows.
- `16 Aug 2020`. [upczww](https://github.com/upczww) added a python wrapper for yolov5.
- `28 May 2020`. arcface LResNet50E-IR model from [deepinsight/insightface](https://github.com/deepinsight/insightface) implemented. We got 333fps on GTX1080.
- `22 May 2020`. A new branch [trt4](https://github.com/wang-xinyu/tensorrtx/tree/trt4) created, which is using TensorRT 4 API. Now the master branch is using TensorRT 7 API. But only `yolov4` has been migrated to TensorRT 7 API for now. The rest will be migrated soon. And a tutorial for `migarating from TensorRT 4 to 7` provided.
## Tutorials
@ -79,6 +79,7 @@ Following models are implemented.
|[hrnet](./hrnet)| hrnet-image-classification and hrnet-semantic-segmentation, pytorch implementation from [HRNet-Image-Classification](https://github.com/HRNet/HRNet-Image-Classification) and [HRNet-Semantic-Segmentation](https://github.com/HRNet/HRNet-Semantic-Segmentation) |
|[psenet](./psenet)| PSENet Text Detection, tensorflow implementation from [liuheng92/tensorflow_PSENet](https://github.com/liuheng92/tensorflow_PSENet) |
|[ibnnet](./ibnnet)| IBN-Net, pytorch implementation from [XingangPan/IBN-Net](https://github.com/XingangPan/IBN-Net), ECCV2018 |
|[unet](./unet)| U-Net, pytorch implementation from [milesial/Pytorch-UNet](https://github.com/milesial/Pytorch-UNet) |
## Model Zoo

2
hrnet/hrnet-image-classification/CMakeLists.txt
View File

@ -13,7 +13,7 @@ include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
find_package(OpenCV)
include_directories(OpenCV_INCLUDE_DIRS)
include_directories(${OpenCV_INCLUDE_DIRS})
add_executable(hrnet ${PROJECT_SOURCE_DIR}/hrnet.cpp)
target_link_libraries(hrnet nvinfer)

2
hrnet/hrnet-semantic-segmentation/CMakeLists.txt
View File

@ -22,7 +22,7 @@ link_directories(/usr/lib/x86_64-linux-gnu/)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
find_package(OpenCV)
include_directories(OpenCV_INCLUDE_DIRS)
include_directories(${OpenCV_INCLUDE_DIRS})
add_executable(hrnetseg ${PROJECT_SOURCE_DIR}/hrnetseg.cpp)
target_link_libraries(hrnetseg nvinfer)

4
unet/CMakeLists.txt
View File

@ -30,5 +30,5 @@ add_definitions(-O2 -pthread)
# opencv library
find_package(OpenCV)
include_directories(OpenCV_INCLUDE_DIRS)
target_link_libraries(unet ${OpenCV_LIBS})
include_directories(${OpenCV_INCLUDE_DIRS})
target_link_libraries(unet ${OpenCV_LIBS})

2
yolov5/calibrator.cpp
View File

@ -3,7 +3,7 @@
#include <fstream>
#include <opencv2/dnn/dnn.hpp>
#include "calibrator.h"
#include "cuda_runtime_api.h"
#include "cuda_utils.h"
#include "utils.h"
Int8EntropyCalibrator2::Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache)

18
yolov5/cuda_utils.h Normal file
View File

@ -0,0 +1,18 @@
#ifndef TRTX_CUDA_UTILS_H_
#define TRTX_CUDA_UTILS_H_
#include <cuda_runtime_api.h>
#ifndef CUDA_CHECK
#define CUDA_CHECK(callstr)\
{\
cudaError_t error_code = callstr;\
if (error_code != cudaSuccess) {\
std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__;\
assert(0);\
}\
}
#endif // CUDA_CHECK
#endif // TRTX_CUDA_UTILS_H_

41
yolov5/utils.h
View File

@ -1,43 +1,9 @@
#ifndef __TRT_UTILS_H_
#define __TRT_UTILS_H_
#ifndef TRTX_YOLOV5_UTILS_H_
#define TRTX_YOLOV5_UTILS_H_
#include <iostream>
#include <vector>
#include <algorithm>
#include <cudnn.h>
#include <dirent.h>
#include <opencv2/opencv.hpp>
#ifndef CUDA_CHECK
#define CUDA_CHECK(callstr) \
{ \
cudaError_t error_code = callstr; \
if (error_code != cudaSuccess) { \
std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \
assert(0); \
} \
}
#endif
namespace Tn
{
template<typename T>
void write(char*& buffer, const T& val)
{
*reinterpret_cast<T*>(buffer) = val;
buffer += sizeof(T);
}
template<typename T>
void read(const char*& buffer, T& val)
{
val = *reinterpret_cast<const T*>(buffer);
buffer += sizeof(T);
}
}
static inline cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) {
int w, h, x, y;
float r_w = input_w / (img.cols*1.0);
@ -82,4 +48,5 @@ static inline int read_files_in_dir(const char *p_dir_name, std::vector<std::str
return 0;
}
#endif
#endif // TRTX_YOLOV5_UTILS_H_

22
yolov5/yololayer.cu
View File

@ -1,6 +1,25 @@
#include <assert.h>
#include <vector>
#include <iostream>
#include "yololayer.h"
#include "utils.h"
#include "cuda_utils.h"
namespace Tn
{
template<typename T>
void write(char*& buffer, const T& val)
{
*reinterpret_cast<T*>(buffer) = val;
buffer += sizeof(T);
}
template<typename T>
void read(const char*& buffer, T& val)
{
val = *reinterpret_cast<const T*>(buffer);
buffer += sizeof(T);
}
}
using namespace Yolo;
@ -311,3 +330,4 @@ namespace nvinfer1
return obj;
}
}

2
yolov5/yolov5.cpp
View File

@ -1,6 +1,6 @@
#include <iostream>
#include <chrono>
#include "cuda_runtime_api.h"
#include "cuda_utils.h"
#include "logging.h"
#include "common.hpp"
#include "utils.h"