yolo目标检测线程池高性能视频tensorrt推理（每秒1400帧）
·
yolov5-tensorrt推理框架实现
logger.h
#ifndef LOGGER_H
#define LOGGER_H
#include <NvInfer.h>
#include <NvInferRuntime.h>
inline const char* severity_string(nvinfer1::ILogger::Severity t)
{
switch (t)
{
case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "internal_error";
case nvinfer1::ILogger::Severity::kERROR: return "error";
case nvinfer1::ILogger::Severity::kWARNING: return "warning";
case nvinfer1::ILogger::Severity::kINFO: return "info";
case nvinfer1::ILogger::Severity::kVERBOSE: return "verbose";
default: return "unknow";
}
}
class TRTLogger : public nvinfer1::ILogger
{
public:
virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override
{
if (severity <= Severity::kINFO)
{
if (severity == Severity::kWARNING)
printf("\033[33m%s: %s\033[0m\n", severity_string(severity), msg);
else if (severity <= Severity::kERROR)
printf("\033[31m%s: %s\033[0m\n", severity_string(severity), msg);
else
printf("%s: %s\n", severity_string(severity), msg);
}
}
} logger;
#endif // LOGGER_H
preprocess.h
#ifndef PREPROCESS_H
#define PREPROCESS_H
#include <opencv2/opencv.hpp>
void LetterBox(const cv::Mat& image, cv::Mat& outImage, const cv::Size& newShape = cv::Size(640, 640), const cv::Scalar& color = cv::Scalar(114, 114, 114));
#endif // PREPROCESS_H
preprocess.cpp
#include "preprocess.h"
void LetterBox(const cv::Mat& image, cv::Mat& outImage, const cv::Size& newShape, const cv::Scalar& color)
{
cv::Size shape = image.size();
float r = std::min((float)newShape.height / (float)shape.height, (float)newShape.width / (float)shape.width);
float ratio[2]{ r, r };
int new_un_pad[2] = { (int)std::round((float)shape.width * r),(int)std::round((float)shape.height * r) };
auto dw = (float)(newShape.width - new_un_pad[0]) / 2;
auto dh = (float)(newShape.height - new_un_pad[1]) / 2;
if (shape.width != new_un_pad[0] && shape.height != new_un_pad[1])
cv::resize(image, outImage, cv::Size(new_un_pad[0], new_un_pad[1]));
else
outImage = image.clone();
int top = int(std::round(dh - 0.1f));
int bottom = int(std::round(dh + 0.1f));
int left = int(std::round(dw - 0.1f));
int right = int(std::round(dw + 0.1f));
cv::Vec4d params;
params[0] = ratio[0];
params[1] = ratio[1];
params[2] = left;
params[3] = top;
cv::copyMakeBorder(outImage, outImage, top, bottom, left, right, cv::BORDER_CONSTANT, color);
}
yolov5.h
#ifndef YOLOV5_H
#define YOLOV5_H
#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <ctime>
#include <cuda_runtime.h>
#include <NvInfer.h>
#include <NvInferRuntime.h>
#include <opencv2/opencv.hpp>
#include "preprocess.h"
#include "postprocess.h"
//#define USE_CUDA
#ifdef USE_CUDA
#include "preprocess.cuh"
#include "decode.cuh"
#endif
class Yolov5
{
public:
Yolov5();
~Yolov5();
int load_model(const std::string& model_path); // 加载模型
int infer(const cv::Mat &image, std::vector<Detection> &detections); // 推理运行模型
private:
int pre_process(const cv::Mat &image); // 图像预处理
int post_process(const cv::Mat &image, std::vector<Detection>& detections); // 后处理
const cv::Size input_size = cv::Size(640, 640);
const int input_numel = 1 * 3 * input_size.width * input_size.height;
const float confidence_threshold = 0.5;
const float score_threshold = 0.25;
const float nms_threshold = 0.45;
const int class_num = 80;
const int output_numprob = 5 + class_num;
const int output_numbox = 3 * (input_size.width / 8 * input_size.height / 8 + input_size.width / 16 * input_size.height / 16 + input_size.width / 32 * input_size.height / 32);
const int output_numel = 1 * output_numprob * output_numbox;
nvinfer1::IRuntime* runtime = nullptr;
nvinfer1::ICudaEngine* engine = nullptr;
nvinfer1::IExecutionContext* execution_context = nullptr;
cudaStream_t stream = nullptr;
float* input_h = nullptr;
float* output_h = nullptr;
float* input_d = nullptr;
float* output_d = nullptr;
float* bindings[2];
#ifdef USE_CUDA
uint8_t* input_host;
float* d2s_host;
float* d2s_device;
float* s2d_host;
float* s2d_device;
float* output_box_host;
float* output_box_device;
const int max_box = 1024;
const int nubox_element = 7;
const int max_input_size = sizeof(float) * 3 * 1024 * 1024;
#endif
};
#endif // YOLOV5_H
yolov5.cpp
#include "yolov5.h"
#include "logger.h"
Yolov5::Yolov5()
{
runtime = nvinfer1::createInferRuntime(logger);
cudaMallocHost(&input_h, sizeof(float) * input_numel);
cudaMallocHost(&output_h, sizeof(float) * output_numel);
cudaMalloc(&input_d, sizeof(float) * input_numel);
cudaMalloc(&output_d, sizeof(float) * output_numel);
bindings[0] = input_d;
bindings[1] = output_d;
#ifdef USE_CUDA
cudaMallocHost(&input_host, max_input_size);
cudaMallocHost(&d2s_host, sizeof(float) * 6);
cudaMalloc(&d2s_device, sizeof(float) * 6);
cudaMallocHost(&s2d_host, sizeof(float) * 6);
cudaMalloc(&s2d_device, sizeof(float) * 6);
cudaMallocHost(&output_box_host, sizeof(float) * (nubox_element * max_box + 1));
cudaMalloc(&output_box_device, sizeof(float) * (nubox_element * max_box + 1));
#endif
}
Yolov5::~Yolov5()
{
cudaStreamDestroy(stream);
cudaFree(input_d);
cudaFree(output_d);
cudaFreeHost(input_h);
cudaFreeHost(output_h);
}
int Yolov5::load_model(const std::string& model_path)
{
std::ifstream in(model_path, std::ios::binary);
if (!in.is_open())
return -1;
in.seekg(0, std::ios::end);
size_t size = in.tellg();
std::vector<unsigned char> engine_data(size);
in.seekg(0, std::ios::beg);
in.read((char*)engine_data.data(), size);
in.close();
engine = runtime->deserializeCudaEngine(engine_data.data(), engine_data.size());
if(engine == nullptr)
return -1;
execution_context = engine->createExecutionContext();
if(execution_context == nullptr)
return -1;
cudaStreamCreate(&stream);
return 0;
}
int Yolov5::pre_process(const cv::Mat &image)
{
#ifdef USE_CUDA
cudaMemcpyAsync(input_host, image.data, sizeof(uint8_t) * 3 * image.cols * image.rows, cudaMemcpyHostToDevice, stream);
preprocess_kernel_img(input_host, image.cols, image.rows, input_d, input_size.width, input_size.height, d2s_host, s2d_host, stream);
cudaMemcpyAsync(d2s_device, d2s_host, sizeof(float) * 6, cudaMemcpyHostToDevice, stream);
cudaMemcpyAsync(s2d_device, s2d_host, sizeof(float) * 6, cudaMemcpyHostToDevice, stream);
#else
cv::Mat letterbox;
LetterBox(image, letterbox, input_size);
//cv::resize(image, letterbox, input_size);
letterbox.convertTo(letterbox, CV_32FC3, 1.0f / 255.0f);
int image_area = letterbox.cols * letterbox.rows;
float* pimage = (float*)letterbox.data;
float* phost_b = input_h + image_area * 0;
float* phost_g = input_h + image_area * 1;
float* phost_r = input_h + image_area * 2;
for (int i = 0; i < image_area; ++i, pimage += 3)
{
*phost_r++ = pimage[0];
*phost_g++ = pimage[1];
*phost_b++ = pimage[2];
}
cudaMemcpyAsync(input_d, input_h, sizeof(float) * input_numel, cudaMemcpyHostToDevice, stream);
#endif
return 0;
}
int Yolov5::infer(const cv::Mat& image, std::vector<Detection>& detections)
{
pre_process(image);
bool success = execution_context->executeV2((void**)bindings);
if(!success)
{
std::cerr << "Failed to run inference" << std::endl;
return -1;
}
#ifndef USE_CUDA
cudaMemcpyAsync(output_h, output_d, sizeof(float) * output_numel, cudaMemcpyDeviceToHost, stream);
cudaStreamSynchronize(stream);
#endif
post_process(image, detections);
return 0;
}
int Yolov5::post_process(const cv::Mat &image, std::vector<Detection>& detections)
{
std::vector<cv::Rect> boxes;
std::vector<float> scores;
std::vector<int> class_ids;
#ifdef USE_CUDA
cudaMemset(output_box_device, 0, sizeof(float) * (nubox_element * max_box + 1));
decode_kernel_invoker(output_d, output_numbox, class_num, score_threshold, d2s_device, output_box_device, max_box, nubox_element, stream);
nms_kernel_invoker(output_box_device, nms_threshold, max_box, nubox_element, stream);
cudaMemcpyAsync(output_box_host, output_box_device, sizeof(float) * (nubox_element * max_box + 1), cudaMemcpyDeviceToHost, stream);
cudaStreamSynchronize(stream);
for (size_t i = 0; i < max_box; i++)
{
if (output_box_host[7 * i + 7])
{
float x1 = output_box_host[7 * i + 1];
float y1 = output_box_host[7 * i + 2];
float x2 = output_box_host[7 * i + 3];
float y2 = output_box_host[7 * i + 4];
boxes.push_back(cv::Rect(x1, y1, x2-x1, y2-y1));
scores.push_back(output_box_host[7 * i + 5]);
class_ids.push_back(output_box_host[7 * i + 6]);
}
}
detections.clear();
detections.resize(boxes.size());
for (int i = 0; i < boxes.size(); i++)
{
detections[i].bbox = boxes[i];
detections[i].score = scores[i];
detections[i].id = class_ids[i];
}
#else
// float x_ratio = float(image.cols) / input_size.width;
// float y_ratio = float(image.rows) / input_size.height;
for (int i = 0; i < output_numbox; ++i)
{
float* ptr = output_h + i * output_numprob;
float obj_score = ptr[4];
if (obj_score < confidence_threshold)
continue;
float* classes_scores = 5 + ptr;
int class_id = std::max_element(classes_scores, classes_scores + class_num) - classes_scores;
float score = classes_scores[class_id] * obj_score;
if (score < score_threshold)
continue;
float x = ptr[0];
float y = ptr[1];
float w = ptr[2];
float h = ptr[3];
int left = int(x - 0.5 * w);
int top = int(y - 0.5 * h);
int width = int(w);
int height = int(h);
cv::Rect box = cv::Rect(left, top, width, height);
scale_boxes(box, input_size, image.size());
boxes.push_back(box);
scores.push_back(score);
class_ids.push_back(class_id);
}
std::vector<int> indices;
nms(boxes, scores, score_threshold, nms_threshold, indices);
detections.clear();
detections.resize(indices.size());
for (int i = 0; i < indices.size(); ++i)
{
int idx = indices[i];
detections[i].bbox = boxes[idx];
detections[i].score = scores[idx];
detections[i].id = class_ids[idx];
}
#endif
return 0;
}
单线程版本
test_yolov5.cpp
#include "yolov5.h"
int main()
{
Yolov5* yolov5 = new Yolov5();
yolov5->load_model("yolov5n_int8.engine");
cv::Mat image = cv::imread("bus.jpg");
std::vector<Detection> detections;
yolov5->infer(image, detections);
std::cout << "detections size: " << detections.size() << std::endl;
yolov5->draw_detections(0, image, detections);
cv::imwrite("result.jpg", image);
cv::VideoCapture cap("bj_full.mp4");
if (!cap.isOpened())
return -1;
int frame_count = 0;
auto start_all = std::chrono::high_resolution_clock::now();
while (true)
{
cap >> image;
if (image.empty())
break;
yolov5->infer(image, detections);
frame_count++;
auto end_all = std::chrono::high_resolution_clock::now();
auto elapsed_all = std::chrono::duration_cast<std::chrono::microseconds>(end_all - start_all).count() / 1000.f;
if (elapsed_all >= 1000)
{
printf("FPS:%f \n", frame_count / (elapsed_all / 1000.0f));
frame_count = 0;
start_all = std::chrono::high_resolution_clock::now();
}
}
cap.release();
return 0;
}
运行./demo输出：
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChane] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +15, now: CPU 0, GPU 17 (MiB)
detections size: 2
FPS:525.963196
FPS:520.842163
FPS:503.353210
FPS:533.499023
......
资源占用：（测试机器为RTX4090+24核cpu）
可以看到，gpu的utils在40%左右，有很大的优化空间。
线程池版本
yolov5_thread_pool.h
#ifndef YOLOV5_THREAD_POOL_H
#define YOLOV5_THREAD_POOL_H
#include "yolov5.h"
#include <iostream>
#include <vector>
#include <queue>
#include <map>
#include <thread>
#include <mutex>
#include <condition_variable>
class Yolov5ThreadPool
{
private:
std::queue<std::pair<int, cv::Mat>> tasks; // <id, img>用来存放任务
std::vector<std::shared_ptr<Yolov5>> yolov5_instances; // 模型实例
std::map<int, std::vector<Detection>> results; // <id, objects>用来存放结果（检测框）
std::map<int, cv::Mat> img_results; // <id, img>用来存放结果（图片）
std::vector<std::thread> threads; // 线程池
std::mutex mtx1;
std::mutex mtx2;
std::condition_variable cv_task, cv_result;
bool stop;
void worker(int id);
public:
Yolov5ThreadPool();
~Yolov5ThreadPool();
int setUp(const std::string &model_path, int num_threads = 12); // 初始化
int submitTask(const cv::Mat &img, int id); // 提交任务
int getTargetResult(std::vector<Detection> &objects, int id); // 获取结果
int getTargetImgResult(cv::Mat &img, int id); // 获取结果（图片）
void stopAll(); // 停止所有线程
};
#endif //YOLOV5_THREAD_POOL_H
yolov5_thread_pool.cpp
#include "yolov5_thread_pool.h"
// 构造函数
Yolov5ThreadPool::Yolov5ThreadPool() { stop = false; }
// 析构函数
Yolov5ThreadPool::~Yolov5ThreadPool()
{
// stop all threads
stop = true;
cv_task.notify_all();
for (auto &thread : threads)
{
if (thread.joinable())
{
thread.join();
}
}
}
// 初始化：加载模型，创建线程，参数：模型路径，线程数量
int Yolov5ThreadPool::setUp(const std::string &model_path, int num_threads)
{
// 遍历线程数量，创建模型实例，放入vector
// 这些线程加载的模型是同一个
for (size_t i = 0; i < num_threads; ++i)
{
std::shared_ptr<Yolov5> yolov5 = std::make_shared<Yolov5>();
yolov5->load_model(model_path.c_str());
yolov5_instances.push_back(yolov5);
}
// 遍历线程数量，创建线程
for (size_t i = 0; i < num_threads; ++i)
{
threads.emplace_back(&Yolov5ThreadPool::worker, this, i);
}
return 0;
}
// 线程函数。参数：线程id
void Yolov5ThreadPool::worker(int id)
{
while (!stop)
{
std::pair<int, cv::Mat> task;
std::shared_ptr<Yolov5> instance = yolov5_instances[id]; // 获取模型实例
{
// 获取任务
std::unique_lock<std::mutex> lock(mtx1);
cv_task.wait(lock, [&] { return !tasks.empty() || stop; });
if (stop)
return;
task = tasks.front();
tasks.pop();
}
// 运行模型
std::vector<Detection> detections;
instance->infer(task.second, detections);
{
// 保存结果
std::lock_guard<std::mutex> lock(mtx2);
results.insert({task.first, detections});
draw_detections(task.second, detections);
img_results.insert({task.first, task.second});
cv_result.notify_one();
}
}
}
// 提交任务，参数：图片，id（帧号）
int Yolov5ThreadPool::submitTask(const cv::Mat &img, int id)
{
// 如果任务队列中的任务数量大于10，等待，避免内存占用过多
while (tasks.size() > 1000)
{
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
{
// 保存任务
std::lock_guard<std::mutex> lock(mtx1);
tasks.push({id, img});
}
cv_task.notify_one();
return 0;
}
// 获取结果，参数：检测框，id（帧号）
int Yolov5ThreadPool::getTargetResult(std::vector<Detection> &objects, int id)
{
// 如果没有结果，等待
while (results.find(id) == results.end())
{
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
std::lock_guard<std::mutex> lock(mtx2);
objects = results[id];
// remove from map
results.erase(id);
return 0;
}
// 获取结果（图片），参数：图片，id（帧号）
int Yolov5ThreadPool::getTargetImgResult(cv::Mat &img, int id)
{
int loop_cnt = 0;
// 如果没有结果，等待
while (img_results.find(id) == img_results.end())
{
// 等待
std::this_thread::sleep_for(std::chrono::milliseconds(1));
loop_cnt++;
if (loop_cnt > 1000)
{
std::cerr << "getTargetImgResult timeout" << std::endl;
return 0;
}
}
std::lock_guard<std::mutex> lock(mtx2);
img = img_results[id];
// remove from map
img_results.erase(id);
return 0;
}
void Yolov5ThreadPool::stopAll()
{
stop = true;
cv_task.notify_all();
}
test_yolov5_thread_pool.cpp
#include "yolov5_thread_pool.h"
#include <chrono>
static int g_frame_start_id = 0; // 读取视频帧的索引
static int g_frame_end_id = 0; // 模型处理完的索引
static Yolov5ThreadPool *g_pool = nullptr;
bool end = false;
void read_stream(const std::string& video_file)
{
cv::VideoCapture cap(video_file);
if (!cap.isOpened())
return;
cv::Mat img;
while (true)
{
cap >> img;
if (img.empty())
{
end = true;
break;
}
g_pool->submitTask(img.clone(), g_frame_start_id++);
}
cap.release();
}
void get_results()
{
auto start_all = std::chrono::high_resolution_clock::now();
int frame_count = 0;
//cv::VideoWriter writer = cv::VideoWriter("result.mp4", cv::VideoWriter::fourcc('m', 'p', '4', 'v'), 30, cv::Size(1280, 720));
while (true)
{
cv::Mat img;
auto ret = g_pool->getTargetImgResult(img, g_frame_end_id++);
if (end)
{
g_pool->stopAll();
break;
}
//cv::imwrite("output/" + std::to_string(g_frame_end_id) + ".jpg", img);
//writer << img;
frame_count++;
auto end_all = std::chrono::high_resolution_clock::now();
auto elapsed_all_2 = std::chrono::duration_cast<std::chrono::microseconds>(end_all - start_all).count() / 1000.f;
if (elapsed_all_2 >= 1000)
{
printf("FPS:%f \n", frame_count / (elapsed_all_2 / 1000.0f));
frame_count = 0;
start_all = std::chrono::high_resolution_clock::now();
}
}
g_pool->stopAll();
}
int main(int argc, char **argv)
{
g_pool = new Yolov5ThreadPool();
g_pool->setUp(argv[1], atoi(argv[2]));
std::thread read_stream_thread(read_stream, "bj_full.mp4");
std::thread result_thread(get_results);
read_stream_thread.join();
result_thread.join();
return 0;
}
CMakeLists.txt
cmake_minimum_required(VERSION 3.20)
project(trt_inference LANGUAGES C CXX CUDA)
set(CMAKE_CXX_STANDARD 14)
find_package(CUDA REQUIRED)
include_directories(${CUDA_INCLUDE_DIRS})
set(OpenCV_DIR /home/tfy/document/HybrIK/cpp/opencv-4.12.0/lib/cmake/opencv4)
find_package(OpenCV REQUIRED)
set(TENSORRT_INCLUDE_DIRS /home/tfy/docker_share/TensorRT-10.6.0.26/include)
set(TENSORRT_LIBRARY_DIRS /home/tfy/docker_share/TensorRT-10.6.0.26/targets/x86_64-linux-gnu/lib)
include_directories(${TENSORRT_INCLUDE_DIRS} ${OpenCV_INCLUDE_DIRS})
link_directories(${TENSORRT_LIBRARY_DIRS})
add_executable(demo test_yolov5.cpp preprocess.cpp postprocess.cpp yolov5.cpp)
add_executable(thread_pool test_yolov5_thread_pool.cpp preprocess.cpp postprocess.cpp yolov5.cpp yolov5_thread_pool.cpp)
target_link_libraries(demo PRIVATE ${OpenCV_LIBS} ${CUDA_LIBRARIES} ${TENSORRT_LIBRARY_DIRS}/libnvinfer.so)
target_link_libraries(thread_pool PRIVATE ${OpenCV_LIBS} ${CUDA_LIBRARIES} ${TENSORRT_LIBRARY_DIRS}/libnvinfer.so)
运行./thread_pool yolov5n_int8.engine 12测试性能，gpu的utils能在70%左右，cpu能跑到近2000%，继续增加线程数时推理帧率可能会下降：
打印输出如下内容，可以看到推理帧率偶尔能达到1000FPS：
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +15, now: CPU 0, GPU 17 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +1, GPU +15, now: CPU 1, GPU 35 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +14, now: CPU 1, GPU 52 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +1, GPU +15, now: CPU 2, GPU 70 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +15, now: CPU 2, GPU 88 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +1, GPU +14, now: CPU 3, GPU 105 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +15, now: CPU 3, GPU 123 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +1, GPU +15, now: CPU 4, GPU 141 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +14, now: CPU 4, GPU 158 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +1, GPU +15, now: CPU 5, GPU 176 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +15, now: CPU 5, GPU 194 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +1, GPU +15, now: CPU 6, GPU 211 (MiB)
FPS:670.522827
FPS:712.649414
FPS:841.653320
FPS:822.008606
FPS:739.883850
FPS:614.572327
FPS:727.774353
FPS:699.419434
FPS:812.266663
FPS:773.891296
FPS:676.168945
FPS:706.710266
FPS:852.323303
FPS:738.334290
FPS:649.812866
FPS:730.089355
FPS:694.656860
FPS:811.991943
FPS:695.474304
FPS:690.040283
FPS:632.997681
FPS:844.942078
FPS:871.944153
FPS:818.565552
FPS:782.801514
FPS:674.829834
FPS:735.527466
FPS:667.501587
FPS:653.889221
FPS:637.249756
FPS:738.310974
FPS:629.966736
FPS:623.615601
FPS:799.024475
FPS:1048.638184
FPS:1083.996826
FPS:742.771179
FPS:703.455505
FPS:800.101257
FPS:1058.399902
nvcodec硬件解码优化
安装ffnvcodec
sudo apt-get update
sudo apt-get install libnvidia-encode-<version> ffmpeg
sudo apt-get install nv-codec-headers
其中version通过下面命令确定:
nvidia-smi | grep "Driver Version"
若输出:
| NVIDIA-SMI 570.211.01 Driver Version: 570.211.01 CUDA Version: 12.8 |
则执行：
sudo apt-get install libnvidia-encode-570 ffmpeg
安装支持CUDA的FFmpeg
LZ下载的是FFmpeg-n6.1.4，运行下面命令编译源码：
./configure --prefix=./install --enable-gpl --enable-nonfree --enable-cuda --enable-cuda-nvcc --enable-cuvid --enable-nvdec --enable-nvenc --disable-debug --enable-shared --enable-static --extra-cfl
make -j8
make install
添加文件nvdec_decode.h
#include <iostream>
#include <chrono>
#include <string>
// CUDA 头文件
#include <cuda.h>
#include <cuda_runtime.h>
// NPP 头文件
#include <npp.h>
// FFmpeg 头文件
extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/hwcontext.h>
#include <libavutil/hwcontext_cuda.h>
#include <libavutil/imgutils.h>
#include <libavutil/pixdesc.h>
}
// OpenCV 头文件
#include <opencv2/opencv.hpp>
/**
* @brief 严格校验 NV12 格式的 CUDA 帧参数（仅需执行一次）
*/
int check_nv12_frame_once(AVFrame* hw_frame, int& w, int& h, int& y_stride, size_t& nv12_size, size_t& bgr_size)
{
if (hw_frame->format != AV_PIX_FMT_CUDA)
{
std::cerr << "错误：hw_frame 不是 CUDA 格式！" << std::endl;
return -1;
}
// 获取固定的宽高/行步长（同一视频所有帧一致）
w = hw_frame->width;
h = hw_frame->height;
y_stride = hw_frame->linesize[0];
// NV12 总大小 = Y(w*h) + UV(w*h/2)
nv12_size = (size_t)w * h * 3 / 2;
// BGR 总大小（固定）
bgr_size = (size_t)w * h * 3;
// 仅一次校验行步长
if (y_stride < w) {
std::cerr << "错误：Y 平面行步长异常！stride=" << y_stride << " width=" << w << std::endl;
return -1;
}
std::cout << "NV12 帧参数校验完成（全局生效）：" << std::endl;
std::cout << " 宽=" << w << " 高=" << h << " Y行步长=" << y_stride << std::endl;
std::cout << " NV12总大小=" << nv12_size << " BGR总大小=" << bgr_size << std::endl;
return 0;
}
/**
* @brief GPU 端 NV12 → BGR 转换（复用显存/参数）
*/
int nv12_to_bgr_cuda(uint8_t* dev_nv12, int w, int h, int y_stride, uint8_t* dev_bgr)
{
// 1. 计算 NV12 各平面指针（考虑行步长对齐）
Npp8u* pY = dev_nv12;
Npp8u* pUV = dev_nv12 + (size_t)y_stride * h;
// 2. 封装为 NPP 要求的指针数组
const Npp8u* const ppSrc[] = {pY, pUV};
// 3. 复用固定的行步长参数
int nSrcStep = y_stride;
int nDstStep = w * 3;
NppiSize oSizeROI = {w, h};
// 4. 调用 NPP 转换函数
NppStatus npp_ret = nppiNV12ToBGR_8u_P2C3R(
ppSrc, // 输入：Y+UV 指针数组
nSrcStep, // 输入行步长（复用全局参数）
dev_bgr, // 输出：复用的BGR显存指针
nDstStep, // 输出行步长（复用）
oSizeROI // 有效图像尺寸（复用）
);
return 0;
}
修改test_yolov5_thread_pool.cpp：
#include "yolov5_thread_pool.h"
#include "global_var.h"
#include <chrono>
#ifdef USE_NVCODEC
#include "nvdec_decode.h"
#endif
static int g_frame_start_id = 0; // 读取视频帧的索引
static int g_frame_end_id = 0; // 模型处理完的索引
static Yolov5ThreadPool *g_pool = nullptr;
bool end = false;
#ifndef USE_NVCODEC
void read_stream(const std::string& video_file)
{
cv::VideoCapture cap(video_file);
if (!cap.isOpened())
return;
cv::Mat img;
while (true)
{
cap >> img;
if (img.empty())
{
end = true;
break;
}
g_pool->submitImgTask(img.clone(), g_frame_start_id++);
}
cap.release();
}
#else
void read_stream_nvcodec(const std::string& video_file)
{
AVFormatContext* fmt_ctx = nullptr;
AVCodecContext* codec_ctx = nullptr;
AVBufferRef* hw_device_ctx = nullptr;
int video_stream_idx = 0;
int ret = 0;
// 打开视频
ret = avformat_open_input(&fmt_ctx, video_file.c_str(), nullptr, nullptr);
if (ret < 0)
{
printf("avformat_open_input failed\n");
return;
}
else{
printf("avformat_open_input success\n");
}
// 查找流信息
ret = avformat_find_stream_info(fmt_ctx, nullptr);
if (ret < 0)
{
printf("avformat_find_stream_info failed\n");
return;
}
else{
printf("avformat_find_stream_info success\n");
}
// 找视频流
video_stream_idx = av_find_best_stream(fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
// 创建 CUDA 设备上下文
ret = av_hwdevice_ctx_create(&hw_device_ctx, AV_HWDEVICE_TYPE_CUDA, nullptr, nullptr, 0);
if (ret < 0)
{
printf("av_hwdevice_ctx_create failed\n");
return;
}
else{
printf("av_hwdevice_ctx_create success\n");
}
// 强制使用 NVDEC 解码器
const AVCodec* codec = avcodec_find_decoder_by_name("h264_cuvid");
if (!codec) {
codec = avcodec_find_decoder_by_name("hevc_cuvid"); // 兼容 HEVC
}
if (!codec) {
std::cerr << "错误：未找到 NVDEC 解码器（h264_cuvid/hevc_cuvid）" << std::endl;
return;
}
else{
printf("avcodec_find_decoder_by_name success\n");
}
// 初始化解码器上下文
codec_ctx = avcodec_alloc_context3(codec);
if (!codec_ctx) {
std::cerr << "avcodec_alloc_context3 failed" << std::endl;
return;
}
else{
printf("avcodec_alloc_context3 success\n");
}
ret = avcodec_parameters_to_context(codec_ctx, fmt_ctx->streams[video_stream_idx]->codecpar);
if (ret < 0)
{
printf("avcodec_parameters_to_context failed\n");
return;
}
else{
printf("avcodec_parameters_to_context success\n");
}
codec_ctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
codec_ctx->thread_count = 1; // GPU 解码无需多线程
ret = avcodec_open2(codec_ctx, codec, nullptr);
if( ret < 0)
{
printf("avcodec_open2 failed\n");
return;
}
else{
printf("avcodec_open2 success\n");
}
// ========== 2. 全局变量：复用帧参数和显存 ==========
AVPacket pkt;
av_init_packet(&pkt);
pkt.data = nullptr;
pkt.size = 0;
AVFrame* hw_frame = av_frame_alloc();
if (!hw_frame) {
printf("av_frame_alloc failed\n");
return;
}
else{
printf("av_frame_alloc success\n");
}
uint8_t* dev_bgr = nullptr; // 全局复用的BGR显存
bool is_param_init = false; // 参数是否已初始化
int frame_count = 0; // 帧计数器
while (av_read_frame(fmt_ctx, &pkt) >= 0) {
if (pkt.stream_index != video_stream_idx) {
av_packet_unref(&pkt);
continue;
}
// 发送数据包
ret = avcodec_send_packet(codec_ctx, &pkt);
if (ret < 0 && ret != AVERROR(EAGAIN)) {
av_packet_unref(&pkt);
continue;
}
if(ret == AVERROR_EOF)
goto cleanup;
// 接收解码帧
while (ret >= 0) {
ret = avcodec_receive_frame(codec_ctx, hw_frame);
if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
break;
frame_count++;
// ========== 4. 仅第一帧执行：参数校验 + 显存分配 ==========
if (!is_param_init) {
// 一次性校验帧参数
ret = check_nv12_frame_once(hw_frame, global_w, global_h, global_y_stride, global_nv12_size, global_bgr_size);
if (ret < 0) {
av_frame_unref(hw_frame);
av_packet_unref(&pkt);
//goto cleanup;
}
// 一次性分配BGR显存（复用至所有帧）
ret = cudaMalloc(&dev_bgr, global_bgr_size);
is_param_init = true; // 标记参数已初始化
}
// ========== 5. 快速校验：防止极端情况（帧尺寸突变） ==========
if (hw_frame->width != global_w || hw_frame->height != global_h || hw_frame->linesize[0] != global_y_stride) {
std::cerr << "错误：帧参数突变！当前帧宽=" << hw_frame->width << " 全局宽=" << global_w << std::endl;
av_frame_unref(hw_frame);
continue;
}
// ========== 6. GPU 转换（复用显存/参数） ==========
ret = nv12_to_bgr_cuda((uint8_t*)hw_frame->data[0], global_w, global_h, global_y_stride, dev_bgr);
if (ret < 0) {
av_frame_unref(hw_frame);
continue;
}
g_pool->submitImgTask(dev_bgr, g_frame_start_id++);
av_frame_unref(hw_frame);
}
av_packet_unref(&pkt);
}
cleanup:
av_frame_free(&hw_frame);
avcodec_free_context(&codec_ctx);
av_buffer_unref(&hw_device_ctx);
avformat_close_input(&fmt_ctx);
exit(0);
}
#endif
void get_results()
{
auto start_all = std::chrono::high_resolution_clock::now();
int frame_count = 0;
//cv::VideoWriter writer = cv::VideoWriter("result.mp4", cv::VideoWriter::fourcc('m', 'p', '4', 'v'), 30, cv::Size(1280, 720));
while (true)
{
cv::Mat img;
std::vector<Detection> detections;
auto ret = g_pool->getTargetResult(detections, g_frame_end_id++);
//auto ret = g_pool->getTargetImgResult(img, g_frame_end_id++);
if (end)
{
g_pool->stopAll();
break;
}
//writer << img;
//std::cout << "frame:" << g_frame_end_id<<" "<<detections.size() << std::endl;
frame_count++;
auto end_all = std::chrono::high_resolution_clock::now();
auto elapsed_all_2 = std::chrono::duration_cast<std::chrono::microseconds>(end_all - start_all).count() / 1000.f;
if (elapsed_all_2 >= 1000)
{
printf("FPS:%f \n", frame_count / (elapsed_all_2 / 1000.0f));
frame_count = 0;
start_all = std::chrono::high_resolution_clock::now();
}
}
g_pool->stopAll();
}
int main(int argc, char **argv)
{
if( argc < 3)
{
printf("Usage: %s <model_path> <video_path> <thread_num>\n", argv[0]);
return -1;
}
g_pool = new Yolov5ThreadPool();
g_pool->setUp(argv[1], atoi(argv[3]));
#ifndef USE_NVCODEC
std::cout << "Using opencv" << std::endl;
std::thread read_stream_thread(read_stream, argv[2]);
#else
std::cout << "Using nvcodec" << std::endl;
std::thread read_stream_thread(read_stream_nvcodec, argv[2]);
#endif
std::thread result_thread(get_results);
read_stream_thread.join();
result_thread.join();
return 0;
}
运行./thread_pool yolov5n_int8.engine bj_full.mp4 2
可以看到，采用nvcodec硬件解码后，gpu的utils能达到近100%，但经过测试增加线程数（>2）时，推理的帧率反而会下降。
打印输出如下内容，推理帧率可以稳定在1400FPS左右：
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +15, now: CPU 0, GPU 17 (MiB)
info: Loaded engine size: 4 MiB
warning: Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.
info: [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +1, GPU +15, now: CPU 1, GPU 35 (MiB)
Using nvcodec
avformat_open_input success
avformat_find_stream_info success
av_hwdevice_ctx_create success
avcodec_find_decoder_by_name success
avcodec_alloc_context3 success
avcodec_parameters_to_context success
[h264_cuvid @ 0x7292ccd83dc0] Invalid pkt_timebase, passing timestamps as-is.
avcodec_open2 success
av_frame_alloc success
NV12 帧参数校验完成（全局生效）：
宽=1280 高=720 Y行步长=1536
NV12总大小=1382400 BGR总大小=2764800
FPS:995.006958
FPS:1402.666138
FPS:1401.745361
FPS:1397.949707
FPS:1396.142700
FPS:1400.215820
FPS:1397.886841
FPS:1395.608521
FPS:1404.317505
FPS:1399.883911
FPS:1401.903687
FPS:1399.732666
FPS:1401.914551
FPS:1401.559082
FPS:1397.706421
FPS:1399.470947
FPS:1401.383423
FPS:1397.094604
FPS:1401.063965
FPS:1400.266235
FPS:1398.696411
FPS:1395.996338
视频推理截图：
完整工程下载链接见：https://download.csdn.net/download/taifyang/92609495
DAMO开发者矩阵，由阿里巴巴达摩院和中国互联网协会联合发起，致力于探讨最前沿的技术趋势与应用成果，搭建高质量的交流与分享平台，推动技术创新与产业应用链接，围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐
9
0- 0
已为社区贡献4条内容
所有评论(0)