preprocess image in GPU for yolov5 (#770)
* Update CMakeLists.txt add nvcc compile option ; add source file : preprocess.cu * Create preprocess.h * Update README.md Add some configurations * Update README.md * Update CMakeLists.txt * Update CMakeLists.txt * Create preprocess.cu * Update yolov5.cpp Preprocess the image in GPU; Use pinned memory to speed up the copy; Only "cudamemcpy" once in function "doInference"; * Update preprocess.h * Update CMakeLists.txt * Update preprocess.cu * Update yolov5.cpp * Update preprocess.cu * Update preprocess.cu * Update preprocess.cu * Update preprocess.cu * Update preprocess.cu * Update preprocess.cu * Update preprocess.cu
This commit is contained in:
liuqi
GitHub
parent
ae329092a1
commit
e7a5f0acbc
@ -23,7 +23,13 @@ link_directories(/usr/local/cuda/lib64)
|
||||
include_directories(/usr/include/x86_64-linux-gnu/)
|
||||
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")
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -g -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
|
||||
|
||||
## configure the compute capability matched with your nvidia graphics card
|
||||
## you need adapt "-gencode" if yours are different
|
||||
## reference for the table for GPU Compute Capability: https://developer.nvidia.com/cuda-gpus#compute
|
||||
set(CUDA_GEN_CODE "-gencode=arch=compute_86,code=sm_86")
|
||||
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -std=c++11 -O0 -Xcompiler -fPIC -g -w ${CUDA_GEN_CODE}")
|
||||
|
||||
cuda_add_library(myplugins SHARED ${PROJECT_SOURCE_DIR}/yololayer.cu)
|
||||
target_link_libraries(myplugins nvinfer cudart)
|
||||
@ -31,7 +37,8 @@ target_link_libraries(myplugins nvinfer cudart)
|
||||
find_package(OpenCV)
|
||||
include_directories(${OpenCV_INCLUDE_DIRS})
|
||||
|
||||
add_executable(yolov5 ${PROJECT_SOURCE_DIR}/calibrator.cpp ${PROJECT_SOURCE_DIR}/yolov5.cpp)
|
||||
cuda_add_executable(yolov5 ${PROJECT_SOURCE_DIR}/calibrator.cpp ${PROJECT_SOURCE_DIR}/yolov5.cpp preprocess.cu)
|
||||
|
||||
target_link_libraries(yolov5 nvinfer)
|
||||
target_link_libraries(yolov5 cudart)
|
||||
target_link_libraries(yolov5 myplugins)
|
||||
@ -41,3 +48,4 @@ if(UNIX)
|
||||
add_definitions(-O2 -pthread)
|
||||
endif(UNIX)
|
||||
|
||||
|
||||
|
||||
@ -20,6 +20,8 @@ Currently, we support yolov5 v1.0, v2.0, v3.0, v3.1, v4.0, v5.0 and v6.0.
|
||||
- Input shape defined in yololayer.h
|
||||
- Number of classes defined in yololayer.h, **DO NOT FORGET TO ADAPT THIS, If using your own model**
|
||||
- INT8/FP16/FP32 can be selected by the macro in yolov5.cpp, **INT8 need more steps, pls follow `How to Run` first and then go the `INT8 Quantization` below**
|
||||
- MAX_IMAGE_INPUT_SIZE_THRESH can be changed in yolov5.cpp ,ensure it exceed the maximum size in the input images!!!! **The default is 3000 x 3000**
|
||||
- The parameter"-gencode" in CMakeLists.txt should be adapted to the compute capability matched with your nvidia graphics card
|
||||
- GPU id can be selected by the macro in yolov5.cpp
|
||||
- NMS thresh in yolov5.cpp
|
||||
- BBox confidence thresh in yolov5.cpp
|
||||
|
||||
126
yolov5/preprocess.cu
Normal file
126
yolov5/preprocess.cu
Normal file
@ -0,0 +1,126 @@
|
||||
#include "preprocess.h"
|
||||
|
||||
|
||||
|
||||
__global__ void warpaffine_kernel(
|
||||
uint8_t* src, int src_line_size, int src_width,
|
||||
int src_height, float* dst, int dst_width,
|
||||
int dst_height, uint8_t const_value_st,
|
||||
AffineMatrix d2s, int edge)
|
||||
{
|
||||
int position = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
if(position >= edge) return;
|
||||
|
||||
float m_x1 = d2s.value[0];
|
||||
float m_y1 = d2s.value[1];
|
||||
float m_z1 = d2s.value[2];
|
||||
float m_x2 = d2s.value[3];
|
||||
float m_y2 = d2s.value[4];
|
||||
float m_z2 = d2s.value[5];
|
||||
|
||||
int dx = position % dst_width;
|
||||
int dy = position / dst_width;
|
||||
float src_x = m_x1 * dx + m_y1 * dy + m_z1 + 0.5f;
|
||||
float src_y = m_x2 * dx + m_y2 * dy + m_z2 + 0.5f;
|
||||
float c0, c1, c2;
|
||||
|
||||
if(src_x <= -1 || src_x >= src_width || src_y <= -1 || src_y >= src_height){
|
||||
// out of range
|
||||
c0 = const_value_st;
|
||||
c1 = const_value_st;
|
||||
c2 = const_value_st;
|
||||
}else{
|
||||
int y_low = floorf(src_y);
|
||||
int x_low = floorf(src_x);
|
||||
int y_high = y_low + 1;
|
||||
int x_high = x_low + 1;
|
||||
|
||||
uint8_t const_value[] = {const_value_st, const_value_st, const_value_st};
|
||||
float ly = src_y - y_low;
|
||||
float lx = src_x - x_low;
|
||||
float hy = 1 - ly;
|
||||
float hx = 1 - lx;
|
||||
float w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
|
||||
uint8_t* v1 = const_value;
|
||||
uint8_t* v2 = const_value;
|
||||
uint8_t* v3 = const_value;
|
||||
uint8_t* v4 = const_value;
|
||||
if(y_low >= 0){
|
||||
if (x_low >= 0)
|
||||
v1 = src + y_low * src_line_size + x_low * 3;
|
||||
|
||||
if (x_high < src_width)
|
||||
v2 = src + y_low * src_line_size + x_high * 3;
|
||||
}
|
||||
|
||||
if(y_high < src_height){
|
||||
if (x_low >= 0)
|
||||
v3 = src + y_high * src_line_size + x_low * 3;
|
||||
|
||||
if (x_high < src_width)
|
||||
v4 = src + y_high * src_line_size + x_high * 3;
|
||||
}
|
||||
|
||||
c0 = w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0];
|
||||
c1 = w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1];
|
||||
c2 = w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2];
|
||||
}
|
||||
|
||||
//bgr to rgb
|
||||
float t = c2;
|
||||
c2 = c0; c0 = t;
|
||||
|
||||
//normalization
|
||||
c0 = c0 / 255.0f;
|
||||
c1 = c1 / 255.0f;
|
||||
c2 = c2 / 255.0f;
|
||||
|
||||
//rgbrgbrgb to rrrgggbbb
|
||||
int area = dst_width * dst_height;
|
||||
float* pdst_c0 = dst + dy * dst_width + dx;
|
||||
float* pdst_c1 = pdst_c0 + area;
|
||||
float* pdst_c2 = pdst_c1 + area;
|
||||
*pdst_c0 = c0;
|
||||
*pdst_c1 = c1;
|
||||
*pdst_c2 = c2;
|
||||
|
||||
}
|
||||
|
||||
|
||||
void preprocess_kernel_img(
|
||||
uint8_t* src, int src_width, int src_height,
|
||||
float* dst, int dst_width, int dst_height,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
AffineMatrix s2d,d2s;
|
||||
float scale = std::min(dst_height / (float)src_height, dst_width / (float)src_width);
|
||||
|
||||
|
||||
s2d.value[0] = scale;
|
||||
s2d.value[1] = 0;
|
||||
s2d.value[2] = -scale * src_width * 0.5 + dst_width * 0.5;
|
||||
s2d.value[3] = 0;
|
||||
s2d.value[4] = scale;
|
||||
s2d.value[5] = -scale * src_height * 0.5 + dst_height * 0.5;
|
||||
|
||||
/*
|
||||
[src_image.x, [dst_image.x,
|
||||
m2x3_s2d * src_image.y, = dst_image.y ]
|
||||
1 ]
|
||||
*/
|
||||
|
||||
cv::Mat m2x3_s2d(2, 3, CV_32F, s2d.value);
|
||||
cv::Mat m2x3_d2s(2, 3, CV_32F, d2s.value);
|
||||
cv::invertAffineTransform(m2x3_s2d, m2x3_d2s);
|
||||
|
||||
memcpy(d2s.value, m2x3_d2s.ptr<float>(0),sizeof(d2s.value));
|
||||
|
||||
int jobs = dst_height * dst_width;
|
||||
int threads = 256;
|
||||
int blocks = ceil(jobs / (float)threads);
|
||||
warpaffine_kernel<<<blocks,threads,0,stream>>>(
|
||||
src, src_width*3, src_width,
|
||||
src_height, dst, dst_width,
|
||||
dst_height, 128, d2s, jobs);
|
||||
|
||||
}
|
||||
22
yolov5/preprocess.h
Normal file
22
yolov5/preprocess.h
Normal file
@ -0,0 +1,22 @@
|
||||
#ifndef PREPROCESS_H
|
||||
#define PREPROCESS_H
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <stdio.h>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include "cuda_utils.h"
|
||||
|
||||
|
||||
using namespace cv;
|
||||
|
||||
struct AffineMatrix{
|
||||
float value[6];
|
||||
};
|
||||
|
||||
|
||||
void preprocess_kernel_img(uint8_t* src, int src_width, int src_height,
|
||||
float* dst, int dst_width, int dst_height,
|
||||
cudaStream_t stream);
|
||||
#endif // PREPROCESS_H
|
||||
@ -6,12 +6,14 @@
|
||||
#include "common.hpp"
|
||||
#include "utils.h"
|
||||
#include "calibrator.h"
|
||||
#include "preprocess.h"
|
||||
|
||||
#define USE_FP16 // set USE_INT8 or USE_FP16 or USE_FP32
|
||||
#define DEVICE 0 // GPU id
|
||||
#define NMS_THRESH 0.4
|
||||
#define CONF_THRESH 0.5
|
||||
#define BATCH_SIZE 1
|
||||
#define MAX_IMAGE_INPUT_SIZE_THRESH 3000 * 3000 // ensure it exceed the maximum size in the input images !
|
||||
|
||||
// stuff we know about the network and the input/output blobs
|
||||
static const int INPUT_H = Yolo::INPUT_H;
|
||||
@ -37,7 +39,7 @@ static int get_depth(int x, float gd) {
|
||||
|
||||
|
||||
|
||||
ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) {
|
||||
ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, nvinfer1::DataType dt, float& gd, float& gw, std::string& wts_name) {
|
||||
INetworkDefinition* network = builder->createNetworkV2(0U);
|
||||
|
||||
// Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
|
||||
@ -124,7 +126,7 @@ ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilder
|
||||
return engine;
|
||||
}
|
||||
|
||||
ICudaEngine* build_engine_p6(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) {
|
||||
ICudaEngine* build_engine_p6(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, nvinfer1::DataType dt, float& gd, float& gw, std::string& wts_name) {
|
||||
INetworkDefinition* network = builder->createNetworkV2(0U);
|
||||
// Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
|
||||
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W });
|
||||
@ -236,9 +238,9 @@ void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream, bool& is_p
|
||||
// Create model to populate the network, then set the outputs and create an engine
|
||||
ICudaEngine *engine = nullptr;
|
||||
if (is_p6) {
|
||||
engine = build_engine_p6(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name);
|
||||
engine = build_engine_p6(maxBatchSize, builder, config, nvinfer1::DataType::kFLOAT, gd, gw, wts_name);
|
||||
} else {
|
||||
engine = build_engine(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name);
|
||||
engine = build_engine(maxBatchSize, builder, config, nvinfer1::DataType::kFLOAT, gd, gw, wts_name);
|
||||
}
|
||||
assert(engine != nullptr);
|
||||
|
||||
@ -251,9 +253,8 @@ void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream, bool& is_p
|
||||
config->destroy();
|
||||
}
|
||||
|
||||
void doInference(IExecutionContext& context, cudaStream_t& stream, void **buffers, float* input, float* output, int batchSize) {
|
||||
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
|
||||
CUDA_CHECK(cudaMemcpyAsync(buffers[0], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
|
||||
void doInference(IExecutionContext& context, cudaStream_t& stream, void **buffers, float* output, int batchSize) {
|
||||
// infer on the batch asynchronously, and DMA output back to host
|
||||
context.enqueue(batchSize, buffers, stream, nullptr);
|
||||
CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
|
||||
cudaStreamSynchronize(stream);
|
||||
@ -350,10 +351,6 @@ int main(int argc, char** argv) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
// prepare input data ---------------------------
|
||||
static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W];
|
||||
//for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
|
||||
// data[i] = 1.0;
|
||||
static float prob[BATCH_SIZE * OUTPUT_SIZE];
|
||||
IRuntime* runtime = createInferRuntime(gLogger);
|
||||
assert(runtime != nullptr);
|
||||
@ -363,7 +360,7 @@ int main(int argc, char** argv) {
|
||||
assert(context != nullptr);
|
||||
delete[] trtModelStream;
|
||||
assert(engine->getNbBindings() == 2);
|
||||
void* buffers[2];
|
||||
float* buffers[2];
|
||||
// In order to bind the buffers, we need to know the names of the input and output tensors.
|
||||
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
|
||||
const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME);
|
||||
@ -371,38 +368,41 @@ int main(int argc, char** argv) {
|
||||
assert(inputIndex == 0);
|
||||
assert(outputIndex == 1);
|
||||
// Create GPU buffers on device
|
||||
CUDA_CHECK(cudaMalloc(&buffers[inputIndex], BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float)));
|
||||
CUDA_CHECK(cudaMalloc(&buffers[outputIndex], BATCH_SIZE * OUTPUT_SIZE * sizeof(float)));
|
||||
CUDA_CHECK(cudaMalloc((void**)&buffers[inputIndex], BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float)));
|
||||
CUDA_CHECK(cudaMalloc((void**)&buffers[outputIndex], BATCH_SIZE * OUTPUT_SIZE * sizeof(float)));
|
||||
|
||||
// Create stream
|
||||
cudaStream_t stream;
|
||||
CUDA_CHECK(cudaStreamCreate(&stream));
|
||||
|
||||
uint8_t* img_host = nullptr;
|
||||
uint8_t* img_device = nullptr;
|
||||
// prepare input data cache in pinned memory
|
||||
CUDA_CHECK(cudaMallocHost((void**)&img_host, MAX_IMAGE_INPUT_SIZE_THRESH * 3));
|
||||
// prepare input data cache in device memory
|
||||
CUDA_CHECK(cudaMalloc((void**)&img_device, MAX_IMAGE_INPUT_SIZE_THRESH * 3));
|
||||
int fcount = 0;
|
||||
for (int f = 0; f < (int)file_names.size(); f++) {
|
||||
fcount++;
|
||||
if (fcount < BATCH_SIZE && f + 1 != (int)file_names.size()) continue;
|
||||
//auto start = std::chrono::system_clock::now();
|
||||
float* buffer_idx = (float*)buffers[inputIndex];
|
||||
for (int b = 0; b < fcount; b++) {
|
||||
cv::Mat img = cv::imread(img_dir + "/" + file_names[f - fcount + 1 + b]);
|
||||
if (img.empty()) continue;
|
||||
cv::Mat pr_img = preprocess_img(img, INPUT_W, INPUT_H); // letterbox BGR to RGB
|
||||
int i = 0;
|
||||
for (int row = 0; row < INPUT_H; ++row) {
|
||||
uchar* uc_pixel = pr_img.data + row * pr_img.step;
|
||||
for (int col = 0; col < INPUT_W; ++col) {
|
||||
data[b * 3 * INPUT_H * INPUT_W + i] = (float)uc_pixel[2] / 255.0;
|
||||
data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0;
|
||||
data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.0;
|
||||
uc_pixel += 3;
|
||||
++i;
|
||||
}
|
||||
}
|
||||
if (img.empty()) continue;
|
||||
size_t size_image = img.cols * img.rows * 3;
|
||||
size_t size_image_dst = INPUT_H * INPUT_W * 3;
|
||||
//copy data to pinned memory
|
||||
memcpy(img_host,img.data,size_image);
|
||||
//copy data to device memory
|
||||
CUDA_CHECK(cudaMemcpyAsync(img_device,img_host,size_image,cudaMemcpyHostToDevice,stream));
|
||||
preprocess_kernel_img(img_device, img.cols, img.rows, buffer_idx, INPUT_W, INPUT_H, stream);
|
||||
buffer_idx += size_image_dst;
|
||||
}
|
||||
|
||||
// Run inference
|
||||
auto start = std::chrono::system_clock::now();
|
||||
doInference(*context, stream, buffers, data, prob, BATCH_SIZE);
|
||||
doInference(*context, stream, (void**)buffers, prob, BATCH_SIZE);
|
||||
auto end = std::chrono::system_clock::now();
|
||||
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
|
||||
std::cout << "inference time:" << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
|
||||
std::vector<std::vector<Yolo::Detection>> batch_res(fcount);
|
||||
for (int b = 0; b < fcount; b++) {
|
||||
auto& res = batch_res[b];
|
||||
@ -410,7 +410,6 @@ int main(int argc, char** argv) {
|
||||
}
|
||||
for (int b = 0; b < fcount; b++) {
|
||||
auto& res = batch_res[b];
|
||||
//std::cout << res.size() << std::endl;
|
||||
cv::Mat img = cv::imread(img_dir + "/" + file_names[f - fcount + 1 + b]);
|
||||
for (size_t j = 0; j < res.size(); j++) {
|
||||
cv::Rect r = get_rect(img, res[j].bbox);
|
||||
@ -420,10 +419,14 @@ int main(int argc, char** argv) {
|
||||
cv::imwrite("_" + file_names[f - fcount + 1 + b], img);
|
||||
}
|
||||
fcount = 0;
|
||||
//auto end = std::chrono::system_clock::now();
|
||||
//std::cout <<"total time:"<< std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
|
||||
}
|
||||
|
||||
// Release stream and buffers
|
||||
cudaStreamDestroy(stream);
|
||||
CUDA_CHECK(cudaFree(img_device));
|
||||
CUDA_CHECK(cudaFreeHost(img_host));
|
||||
CUDA_CHECK(cudaFree(buffers[inputIndex]));
|
||||
CUDA_CHECK(cudaFree(buffers[outputIndex]));
|
||||
// Destroy the engine
|
||||
@ -431,6 +434,7 @@ int main(int argc, char** argv) {
|
||||
engine->destroy();
|
||||
runtime->destroy();
|
||||
|
||||
|
||||
// Print histogram of the output distribution
|
||||
//std::cout << "\nOutput:\n\n";
|
||||
//for (unsigned int i = 0; i < OUTPUT_SIZE; i++)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user