只依赖Tensorrt和opencv的yolov5源代码
simple_yolo.hpp
#ifndef SIMPLE_YOLO_HPP
#define SIMPLE_YOLO_HPP
/*
简单的yolo接口，容易集成但是高性能
*/
#include 这是一个简单的yolo接口的头文件，提供了一些用于创建和使用yolo模型的函数和结构体。
Type和Mode是用于指定yolo模型类型和推理模式的枚举类型。
Box结构体表示检测到的物体的边界框、置信度和类别标签。
Infer类是一个抽象类，它提供了一个commit函数，用于执行单个图像的推理，以及一个commits函数，用于执行多个图像的推理。
trt_version、type_name、mode_string和set_device是一些辅助函数，用于获取TensorRT版本、yolo模型类型和推理模式的名称，并设置使用的GPU设备。
compile函数用于将yolo模型从ONNX格式编译为TensorRT格式，并将结果保存到文件中。
create_infer函数用于创建一个推理器，它从文件中加载TensorRT引擎，并使用指定的GPU设备执行推理。
这个头文件提供了一个简单的yolo接口，用于创建和使用yolo模型。
Type和Mode是用于指定yolo模型类型和推理模式的枚举类型。Type包括V5和X两种类型，Mode包括FP32、FP16和INT8三种模式。
Box结构体表示检测到的物体的边界框、置信度和类别标签。它包括一个默认构造函数和一个带有参数的构造函数，用于初始化边界框的位置、置信度和类别标签。
Infer类是一个抽象类，它提供了一个commit函数和一个commits函数。commit函数用于执行单个图像的推理，并返回包含检测到的物体信息的future对象。commits函数用于执行多个图像的推理，并返回包含检测到的物体信息的多个future对象。
trt_version、type_name、mode_string和set_device是一些辅助函数。trt_version返回TensorRT版本号。type_name返回指定yolo模型类型的名称。mode_string返回指定推理模式的名称。set_device用于设置使用的GPU设备。
compile函数用于将yolo模型从ONNX格式编译为TensorRT格式，并将结果保存到文件中。它包括许多参数，如模型类型、推理模式、最大批处理大小和最大工作空间大小等。如果模式为INT8，则需要提供用于标定的图像文件夹路径以及熵缓存文件路径。
create_infer函数用于创建一个推理器，它从文件中加载TensorRT引擎，并使用指定的GPU设备执行推理。它还包括一些参数，如置信度阈值和NMS阈值等。
总的来说，这个头文件提供了一个简单易用的接口，用于创建和使用yolo模型进行对象检测。
simple_yolo.cu
#include "simple_yolo.hpp"
#include () + offset_location, src, copyed_bytes, cudaMemcpyDeviceToDevice, stream_));
}
else{
checkCudaRuntime(cudaMemcpyAsync(gpu<unsigned char>() + offset_location, src, copyed_bytes, cudaMemcpyDeviceToDevice, stream_));
}
}else if(head_ == DataHead::Host){
AutoDevice auto_device_exchange(this->device());
checkCudaRuntime(cudaMemcpyAsync(cpu<unsigned char>() + offset_location, src, copyed_bytes, cudaMemcpyDeviceToHost, stream_));
}else{
INFOE("Unsupport head type %d", head_);
}
return *this;
}
Tensor& Tensor::release() {
data_->release_all();
shape_.clear();
bytes_ = 0;
head_ = DataHead::Init;
return *this;
}
bool Tensor::empty() const{
return data_->cpu() == nullptr && data_->gpu() == nullptr;
}
int Tensor::count(int start_axis) const {
if(start_axis >= 0 && start_axis < shape_.size()){
int size = 1;
for (int i = start_axis; i < shape_.size(); ++i)
size *= shape_[i];
return size;
}else{
return 0;
}
}
Tensor& Tensor::resize(const std::vector<int>& dims) {
return resize(dims.size(), dims.data());
}
int Tensor::numel() const{
int value = shape_.empty() ? 0 : 1;
for(int i = 0; i < shape_.size(); ++i){
value *= shape_[i];
}
return value;
}
Tensor& Tensor::resize_single_dim(int idim, int size){
Assert(idim >= 0 && idim < shape_.size());
auto new_shape = shape_;
new_shape[idim] = size;
return resize(new_shape);
}
Tensor& Tensor::resize(int ndims, const int* dims) {
vector<int> setup_dims(ndims);
for(int i = 0; i < ndims; ++i){
int dim = dims[i];
if(dim == -1){
Assert(ndims == shape_.size());
dim = shape_[i];
}
setup_dims[i] = dim;
}
this->shape_ = setup_dims;
this->adajust_memory_by_update_dims_or_type();
this->compute_shape_string();
return *this;
}
Tensor& Tensor::adajust_memory_by_update_dims_or_type(){
int needed_size = this->numel() * element_size();
if(needed_size > this->bytes_){
head_ = DataHead::Init;
}
this->bytes_ = needed_size;
return *this;
}
Tensor& Tensor::synchronize(){
AutoDevice auto_device_exchange(this->device());
checkCudaRuntime(cudaStreamSynchronize(stream_));
return *this;
}
Tensor& Tensor::to_gpu(bool copy) {
if (head_ == DataHead::Device)
return *this;
head_ = DataHead::Device;
data_->gpu(bytes_);
if (copy && data_->cpu() != nullptr) {
AutoDevice auto_device_exchange(this->device());
checkCudaRuntime(cudaMemcpyAsync(data_->gpu(), data_->cpu(), bytes_, cudaMemcpyHostToDevice, stream_));
}
return *this;
}
Tensor& Tensor::to_cpu(bool copy) {
if (head_ == DataHead::Host)
return *this;
head_ = DataHead::Host;
data_->cpu(bytes_);
if (copy && data_->gpu() != nullptr) {
AutoDevice auto_device_exchange(this->device());
checkCudaRuntime(cudaMemcpyAsync(data_->cpu(), data_->gpu(), bytes_, cudaMemcpyDeviceToHost, stream_));
checkCudaRuntime(cudaStreamSynchronize(stream_));
}
return *this;
}
int Tensor::offset_array(size_t size, const int* index_array) const{
Assert(size <= shape_.size());
int value = 0;
for(int i = 0; i < shape_.size(); ++i){
if(i < size)
value += index_array[i];
if(i + 1 < shape_.size())
value *= shape_[i+1];
}
return value;
}
int Tensor::offset_array(const std::vector<int>& index_array) const{
return offset_array(index_array.size(), index_array.data());
}
bool Tensor::save_to_file(const std::string& file) const{
if(empty()) return false;
FILE* f = fopen(file.c_str(), "wb");
if(f == nullptr) return false;
int ndims = this->ndims();
int dtype_ = 0;
unsigned int head[3] = {0xFCCFE2E2, ndims, static_cast<unsigned int>(dtype_)};
fwrite(head, 1, sizeof(head), f);
fwrite(shape_.data(), 1, sizeof(shape_[0]) * shape_.size(), f);
fwrite(cpu(), 1, bytes_, f);
fclose(f);
return true;
}
/class TRTInferImpl
class Logger : public ILogger {
public:
virtual void log(Severity severity, const char* msg) noexcept override {
if (severity == Severity::kINTERNAL_ERROR) {
INFOE("NVInfer INTERNAL_ERROR: %s", msg);
abort();
}else if (severity == Severity::kERROR) {
INFOE("NVInfer: %s", msg);
}
else if (severity == Severity::kWARNING) {
INFOW("NVInfer: %s", msg);
}
else if (severity == Severity::kINFO) {
INFOD("NVInfer: %s", msg);
}
else {
INFOD("%s", msg);
}
}
};
static Logger gLogger;
template<typename _T>
static void destroy_nvidia_pointer(_T* ptr) {
if (ptr) ptr->destroy();
}
class EngineContext {
public:
virtual ~EngineContext() { destroy(); }
void set_stream(cudaStream_t stream){
if(owner_stream_){
if (stream_) {cudaStreamDestroy(stream_);}
owner_stream_ = false;
}
stream_ = stream;
}
bool build_model(const void* pdata, size_t size) {
destroy();
if(pdata == nullptr || size == 0)
return false;
owner_stream_ = true;
checkCudaRuntime(cudaStreamCreate(&stream_));
if(stream_ == nullptr)
return false;
runtime_ = shared_ptr<IRuntime>(createInferRuntime(gLogger), destroy_nvidia_pointer<IRuntime>);
if (runtime_ == nullptr)
return false;
engine_ = shared_ptr<ICudaEngine>(runtime_->deserializeCudaEngine(pdata, size, nullptr), destroy_nvidia_pointer<ICudaEngine>);
if (engine_ == nullptr)
return false;
//runtime_->setDLACore(0);
context_ = shared_ptr<IExecutionContext>(engine_->createExecutionContext(), destroy_nvidia_pointer<IExecutionContext>);
return context_ != nullptr;
}
private:
void destroy() {
context_.reset();
engine_.reset();
runtime_.reset();
if(owner_stream_){
if (stream_) {cudaStreamDestroy(stream_);}
}
stream_ = nullptr;
}
public:
cudaStream_t stream_ = nullptr;
bool owner_stream_ = false;
shared_ptr<IExecutionContext> context_;
shared_ptr<ICudaEngine> engine_;
shared_ptr<IRuntime> runtime_ = nullptr;
};
class TRTInferImpl{
public:
virtual ~TRTInferImpl();
bool load(const std::string& file);
bool load_from_memory(const void* pdata, size_t size);
void destroy();
void forward(bool sync);
int get_max_batch_size();
cudaStream_t get_stream();
void set_stream(cudaStream_t stream);
void synchronize();
size_t get_device_memory_size();
std::shared_ptr<MixMemory> get_workspace();
std::shared_ptr<Tensor> input(int index = 0);
std::string get_input_name(int index = 0);
std::shared_ptr<Tensor> output(int index = 0);
std::string get_output_name(int index = 0);
std::shared_ptr<Tensor> tensor(const std::string& name);
bool is_output_name(const std::string& name);
bool is_input_name(const std::string& name);
void set_input (int index, std::shared_ptr<Tensor> tensor);
void set_output(int index, std::shared_ptr<Tensor> tensor);
std::shared_ptr<std::vector<uint8_t>> serial_engine();
void print();
int num_output();
int num_input();
int device();
private:
void build_engine_input_and_outputs_mapper();
private:
std::vector<std::shared_ptr<Tensor>> inputs_;
std::vector<std::shared_ptr<Tensor>> outputs_;
std::vector<int> inputs_map_to_ordered_index_;
std::vector<int> outputs_map_to_ordered_index_;
std::vector<std::string> inputs_name_;
std::vector<std::string> outputs_name_;
std::vector<std::shared_ptr<Tensor>> orderdBlobs_;
std::map<std::string, int> blobsNameMapper_;
std::shared_ptr<EngineContext> context_;
std::vector<void*> bindingsPtr_;
std::shared_ptr<MixMemory> workspace_;
int device_ = 0;
};
TRTInferImpl::~TRTInferImpl(){
destroy();
}
void TRTInferImpl::destroy() {
int old_device = 0;
checkCudaRuntime(cudaGetDevice(&old_device));
checkCudaRuntime(cudaSetDevice(device_));
this->context_.reset();
this->blobsNameMapper_.clear();
this->outputs_.clear();
this->inputs_.clear();
this->inputs_name_.clear();
this->outputs_name_.clear();
checkCudaRuntime(cudaSetDevice(old_device));
}
void TRTInferImpl::print(){
if(!context_){
INFOW("Infer print, nullptr.");
return;
}
INFO("Infer %p detail", this);
INFO("\tMax Batch Size: %d", this->get_max_batch_size());
INFO("\tInputs: %d", inputs_.size());
for(int i = 0; i < inputs_.size(); ++i){
auto& tensor = inputs_[i];
auto& name = inputs_name_[i];
INFO("\t\t%d.%s : shape {%s}", i, name.c_str(), tensor->shape_string());
}
INFO("\tOutputs: %d", outputs_.size());
for(int i = 0; i < outputs_.size(); ++i){
auto& tensor = outputs_[i];
auto& name = outputs_name_[i];
INFO("\t\t%d.%s : shape {%s}", i, name.c_str(), tensor->shape_string());
}
}
std::shared_ptr<std::vector<uint8_t>> TRTInferImpl::serial_engine() {
auto memory = this->context_->engine_->serialize();
auto output = make_shared<std::vector<uint8_t>>((uint8_t*)memory->data(), (uint8_t*)memory->data()+memory->size());
memory->destroy();
return output;
}
bool TRTInferImpl::load_from_memory(const void* pdata, size_t size) {
if (pdata == nullptr || size == 0)
return false;
context_.reset(new EngineContext());
//build model
if (!context_->build_model(pdata, size)) {
context_.reset();
return false;
}
workspace_.reset(new MixMemory());
cudaGetDevice(&device_);
build_engine_input_and_outputs_mapper();
return true;
}
static std::vector<uint8_t> load_file(const string& file){
ifstream in(file, ios::in | ios::binary);
if (!in.is_open())
return {};
in.seekg(0, ios::end);
size_t length = in.tellg();
std::vector<uint8_t> data;
if (length > 0){
in.seekg(0, ios::beg);
data.resize(length);
in.read((char*)&data[0], length);
}
in.close();
return data;
}
bool TRTInferImpl::load(const std::string& file) {
auto data = load_file(file);
if (data.empty())
return false;
context_.reset(new EngineContext());
//build model
if (!context_->build_model(data.data(), data.size())) {
context_.reset();
return false;
}
workspace_.reset(new MixMemory());
cudaGetDevice(&device_);
build_engine_input_and_outputs_mapper();
return true;
}
size_t TRTInferImpl::get_device_memory_size() {
EngineContext* context = (EngineContext*)this->context_.get();
return context->context_->getEngine().getDeviceMemorySize();
}
void TRTInferImpl::build_engine_input_and_outputs_mapper() {
EngineContext* context = (EngineContext*)this->context_.get();
int nbBindings = context->engine_->getNbBindings();
int max_batchsize = context->engine_->getMaxBatchSize();
inputs_.clear();
inputs_name_.clear();
outputs_.clear();
outputs_name_.clear();
orderdBlobs_.clear();
bindingsPtr_.clear();
blobsNameMapper_.clear();
for (int i = 0; i < nbBindings; ++i) {
auto dims = context->engine_->getBindingDimensions(i);
auto type = context->engine_->getBindingDataType(i);
const char* bindingName = context->engine_->getBindingName(i);
dims.d[0] = max_batchsize;
auto newTensor = make_shared<Tensor>(dims.nbDims, dims.d);
newTensor->set_stream(this->context_->stream_);
newTensor->set_workspace(this->workspace_);
if (context->engine_->bindingIsInput(i)) {
//if is input
inputs_.push_back(newTensor);
inputs_name_.push_back(bindingName);
inputs_map_to_ordered_index_.push_back(orderdBlobs_.size());
}
else {
//if is output
outputs_.push_back(newTensor);
outputs_name_.push_back(bindingName);
outputs_map_to_ordered_index_.push_back(orderdBlobs_.size());
}
blobsNameMapper_[bindingName] = i;
orderdBlobs_.push_back(newTensor);
}
bindingsPtr_.resize(orderdBlobs_.size());
}
void TRTInferImpl::set_stream(cudaStream_t stream){
this->context_->set_stream(stream);
for(auto& t : orderdBlobs_)
t->set_stream(stream);
}
cudaStream_t TRTInferImpl::get_stream() {
return this->context_->stream_;
}
int TRTInferImpl::device() {
return device_;
}
void TRTInferImpl::synchronize() {
checkCudaRuntime(cudaStreamSynchronize(context_->stream_));
}
bool TRTInferImpl::is_output_name(const std::string& name){
return std::find(outputs_name_.begin(), outputs_name_.end(), name) != outputs_name_.end();
}
bool TRTInferImpl::is_input_name(const std::string& name){
return std::find(inputs_name_.begin(), inputs_name_.end(), name) != inputs_name_.end();
}
void TRTInferImpl::forward(bool sync) {
EngineContext* context = (EngineContext*)context_.get();
int inputBatchSize = inputs_[0]->size(0);
for(int i = 0; i < context->engine_->getNbBindings(); ++i){
auto dims = context->engine_->getBindingDimensions(i);
auto type = context->engine_->getBindingDataType(i);
dims.d[0] = inputBatchSize;
if(context->engine_->bindingIsInput(i)){
context->context_->setBindingDimensions(i, dims);
}
}
for (int i = 0; i < outputs_.size(); ++i) {
outputs_[i]->resize_single_dim(0, inputBatchSize);
outputs_[i]->to_gpu(false);
}
for (int i = 0; i < orderdBlobs_.size(); ++i)
bindingsPtr_[i] = orderdBlobs_[i]->gpu();
void** bindingsptr = bindingsPtr_.data();
//bool execute_result = context->context_->enqueue(inputBatchSize, bindingsptr, context->stream_, nullptr);
bool execute_result = context->context_->enqueueV2(bindingsptr, context->stream_, nullptr);
if(!execute_result){
auto code = cudaGetLastError();
INFOF("execute fail, code %d[%s], message %s", code, cudaGetErrorName(code), cudaGetErrorString(code));
}
if (sync) {
synchronize();
}
}
std::shared_ptr<MixMemory> TRTInferImpl::get_workspace() {
return workspace_;
}
int TRTInferImpl::num_input() {
return this->inputs_.size();
}
int TRTInferImpl::num_output() {
return this->outputs_.size();
}
void TRTInferImpl::set_input (int index, std::shared_ptr<Tensor> tensor){
Assert(index >= 0 && index < inputs_.size());
this->inputs_[index] = tensor;
int order_index = inputs_map_to_ordered_index_[index];
this->orderdBlobs_[order_index] = tensor;
}
void TRTInferImpl::set_output(int index, std::shared_ptr<Tensor> tensor){
Assert(index >= 0 && index < outputs_.size());
this->outputs_[index] = tensor;
int order_index = outputs_map_to_ordered_index_[index];
this->orderdBlobs_[order_index] = tensor;
}
std::shared_ptr<Tensor> TRTInferImpl::input(int index) {
Assert(index >= 0 && index < inputs_name_.size());
return this->inputs_[index];
}
std::string TRTInferImpl::get_input_name(int index){
Assert(index >= 0 && index < inputs_name_.size());
return inputs_name_[index];
}
std::shared_ptr<Tensor> TRTInferImpl::output(int index) {
Assert(index >= 0 && index < outputs_.size());
return outputs_[index];
}
std::string TRTInferImpl::get_output_name(int index){
Assert(index >= 0 && index < outputs_name_.size());
return outputs_name_[index];
}
int TRTInferImpl::get_max_batch_size() {
Assert(this->context_ != nullptr);
return this->context_->engine_->getMaxBatchSize();
}
std::shared_ptr<Tensor> TRTInferImpl::tensor(const std::string& name) {
Assert(this->blobsNameMapper_.find(name) != this->blobsNameMapper_.end());
return orderdBlobs_[blobsNameMapper_[name]];
}
std::shared_ptr<TRTInferImpl> load_infer(const string& file) {
std::shared_ptr<TRTInferImpl> infer(new TRTInferImpl());
if (!infer->load(file))
infer.reset();
return infer;
}
//class MonopolyAllocator//
/* 独占分配器
通过对tensor做独占管理，具有max_batch * 2个tensor，通过query获取一个
当推理结束后，该tensor释放使用权，即可交给下一个图像使用，内存实现复用
*/
template<class _ItemType>
class MonopolyAllocator{
public:
class MonopolyData{
public:
std::shared_ptr<_ItemType>& data(){ return data_; }
void release(){manager_->release_one(this);}
private:
MonopolyData(MonopolyAllocator* pmanager){manager_ = pmanager;}
private:
friend class MonopolyAllocator;
MonopolyAllocator* manager_ = nullptr;
std::shared_ptr<_ItemType> data_;
bool available_ = true;
};
typedef std::shared_ptr<MonopolyData> MonopolyDataPointer;
MonopolyAllocator(int size){
capacity_ = size;
num_available_ = size;
datas_.resize(size);
for(int i = 0; i < size; ++i)
datas_[i] = std::shared_ptr<MonopolyData>(new MonopolyData(this));
}
virtual ~MonopolyAllocator(){
run_ = false;
cv_.notify_all();
std::unique_lock<std::mutex> l(lock_);
cv_exit_.wait(l, [&](){
return num_wait_thread_ == 0;
});
}
MonopolyDataPointer query(int timeout = 10000){
std::unique_lock<std::mutex> l(lock_);
if(!run_) return nullptr;
if(num_available_ == 0){
num_wait_thread_++;
auto state = cv_.wait_for(l, std::chrono::milliseconds(timeout), [&](){
return num_available_ > 0 || !run_;
});
num_wait_thread_--;
cv_exit_.notify_one();
// timeout, no available, exit program
if(!state || num_available_ == 0 || !run_)
return nullptr;
}
auto item = std::find_if(datas_.begin(), datas_.end(), [](MonopolyDataPointer& item){return item->available_;});
if(item == datas_.end())
return nullptr;
(*item)->available_ = false;
num_available_--;
return *item;
}
int num_available(){
return num_available_;
}
int capacity(){
return capacity_;
}
private:
void release_one(MonopolyData* prq){
std::unique_lock<std::mutex> l(lock_);
if(!prq->available_){
prq->available_ = true;
num_available_++;
cv_.notify_one();
}
}
private:
std::mutex lock_;
std::condition_variable cv_;
std::condition_variable cv_exit_;
std::vector<MonopolyDataPointer> datas_;
int capacity_ = 0;
volatile int num_available_ = 0;
volatile int num_wait_thread_ = 0;
volatile bool run_ = true;
};
/class ThreadSafedAsyncInfer/
/* 异步线程安全的推理器
通过异步线程启动，使得调用方允许任意线程调用把图像做输入，并通过future来获取异步结果
*/
template<class Input, class Output, class StartParam=std::tuple<std::string, int>, class JobAdditional=int>
class ThreadSafedAsyncInfer{
public:
struct Job{
Input input;
Output output;
JobAdditional additional;
MonopolyAllocator<Tensor>::MonopolyDataPointer mono_tensor;
std::shared_ptr<std::promise<Output>> pro;
};
virtual ~ThreadSafedAsyncInfer(){
stop();
}
void stop(){
run_ = false;
cond_.notify_all();
// cleanup jobs
{
std::unique_lock<std::mutex> l(jobs_lock_);
while(!jobs_.empty()){
auto& item = jobs_.front();
if(item.pro)
item.pro->set_value(Output());
jobs_.pop();
}
};
if(worker_){
worker_->join();
worker_.reset();
}
}
bool startup(const StartParam& param){
run_ = true;
std::promise<bool> pro;
start_param_ = param;
worker_ = std::make_shared<std::thread>(&ThreadSafedAsyncInfer::worker, this, std::ref(pro));
return pro.get_future().get();
}
virtual std::shared_future<Output> commit(const Input& input){
Job job;
job.pro = std::make_shared<std::promise<Output>>();
if(!preprocess(job, input)){
job.pro->set_value(Output());
return job.pro->get_future();
}
///
{
std::unique_lock<std::mutex> l(jobs_lock_);
jobs_.push(job);
};
cond_.notify_one();
return job.pro->get_future();
}
virtual std::vector<std::shared_future<Output>> commits(const std::vector<Input>& inputs){
int batch_size = std::min((int)inputs.size(), this->tensor_allocator_->capacity());
std::vector<Job> jobs(inputs.size());
std::vector<std::shared_future<Output>> results(inputs.size());
int nepoch = (inputs.size() + batch_size - 1) / batch_size;
for(int epoch = 0; epoch < nepoch; ++epoch){
int begin = epoch * batch_size;
int end = std::min((int)inputs.size(), begin + batch_size);
for(int i = begin; i < end; ++i){
Job& job = jobs[i];
job.pro = std::make_shared<std::promise<Output>>();
if(!preprocess(job, inputs[i])){
job.pro->set_value(Output());
}
results[i] = job.pro->get_future();
}
///
{
std::unique_lock<std::mutex> l(jobs_lock_);
for(int i = begin; i < end; ++i){
jobs_.emplace(std::move(jobs[i]));
};
}
cond_.notify_one();
}
return results;
}
protected:
virtual void worker(std::promise<bool>& result) = 0;
virtual bool preprocess(Job& job, const Input& input) = 0;
virtual bool get_jobs_and_wait(std::vector<Job>& fetch_jobs, int max_size){
std::unique_lock<std::mutex> l(jobs_lock_);
cond_.wait(l, [&](){
return !run_ || !jobs_.empty();
});
if(!run_) return false;
fetch_jobs.clear();
for(int i = 0; i < max_size && !jobs_.empty(); ++i){
fetch_jobs.emplace_back(std::move(jobs_.front()));
jobs_.pop();
}
return true;
}
virtual bool get_job_and_wait(Job& fetch_job){
std::unique_lock<std::mutex> l(jobs_lock_);
cond_.wait(l, [&](){
return !run_ || !jobs_.empty();
});
if(!run_) return false;
fetch_job = std::move(jobs_.front());
jobs_.pop();
return true;
}
protected:
StartParam start_param_;
std::atomic<bool> run_;
std::mutex jobs_lock_;
std::queue<Job> jobs_;
std::shared_ptr<std::thread> worker_;
std::condition_variable cond_;
std::shared_ptr<MonopolyAllocator<Tensor>> tensor_allocator_;
};
///class YoloTRTInferImpl//
/* Yolo的具体实现
通过上述类的特性，实现预处理的计算重叠、异步垮线程调用，最终拼接为多个图为一个batch进行推理。最大化的利用
显卡性能，实现高性能高可用好用的yolo推理
*/
const char* type_name(Type type){
switch(type){
case Type::V5: return "YoloV5";
case Type::X: return "YoloX";
default: return "Unknow";
}
}
struct AffineMatrix{
float i2d[6]; // image to dst(network), 2x3 matrix
float d2i[6]; // dst to image, 2x3 matrix
void compute(const cv::Size& from, const cv::Size& to){
float scale_x = to.width / (float)from.width;
float scale_y = to.height / (float)from.height;
float scale = std::min(scale_x, scale_y);
i2d[0] = scale; i2d[1] = 0; i2d[2] = -scale * from.width * 0.5 + to.width * 0.5 + scale * 0.5 - 0.5;
i2d[3] = 0; i2d[4] = scale; i2d[5] = -scale * from.height * 0.5 + to.height * 0.5 + scale * 0.5 - 0.5;
cv::Mat m2x3_i2d(2, 3, CV_32F, i2d);
cv::Mat m2x3_d2i(2, 3, CV_32F, d2i);
cv::invertAffineTransform(m2x3_i2d, m2x3_d2i);
}
cv::Mat i2d_mat(){
return cv::Mat(2, 3, CV_32F, i2d);
}
};
using ThreadSafedAsyncInferImpl = ThreadSafedAsyncInfer
<
cv::Mat, // input
BoxArray, // output
tuple<string, int>, // start param
AffineMatrix // additional
>;
class YoloTRTInferImpl : public Infer, public ThreadSafedAsyncInferImpl{
public:
/** 要求在TRTInferImpl里面执行stop，而不是在基类执行stop **/
virtual ~YoloTRTInferImpl(){
stop();
}
virtual bool startup(const string& file, Type type, int gpuid, float confidence_threshold, float nms_threshold){
if(type == Type::V5){
normalize_ = Norm::alpha_beta(1 / 255.0f, 0.0f, ChannelType::SwapRB);
}else if(type == Type::X){
//float mean[] = {0.485, 0.456, 0.406};
//float std[] = {0.229, 0.224, 0.225};
//normalize_ = Norm::mean_std(mean, std, 1/255.0f, ChannelType::Invert);
normalize_ = Norm::None();
}else{
INFOE("Unsupport type %d", type);
}
confidence_threshold_ = confidence_threshold;
nms_threshold_ = nms_threshold;
return ThreadSafedAsyncInferImpl::startup(make_tuple(file, gpuid));
}
virtual void worker(promise<bool>& result) override{
string file = get<0>(start_param_);
int gpuid = get<1>(start_param_);
set_device(gpuid);
auto engine = load_infer(file);
if(engine == nullptr){
INFOE("Engine %s load failed", file.c_str());
result.set_value(false);
return;
}
engine->print();
const int MAX_IMAGE_BBOX = 1024;
const int NUM_BOX_ELEMENT = 7; // left, top, right, bottom, confidence, class, keepflag
Tensor affin_matrix_device;
Tensor output_array_device;
int max_batch_size = engine->get_max_batch_size();
auto input = engine->tensor("images");
auto output = engine->tensor("output");
int num_classes = output->size(2) - 5;
input_width_ = input->size(3);
input_height_ = input->size(2);
tensor_allocator_ = make_shared<MonopolyAllocator<Tensor>>(max_batch_size * 2);
stream_ = engine->get_stream();
gpu_ = gpuid;
result.set_value(true);
input->resize_single_dim(0, max_batch_size).to_gpu();
affin_matrix_device.set_stream(stream_);
// 这里8个值的目的是保证 8 * sizeof(float) % 32 == 0
affin_matrix_device.resize(max_batch_size, 8).to_gpu();
// 这里的 1 + MAX_IMAGE_BBOX结构是，counter + bboxes ...
output_array_device.resize(max_batch_size, 1 + MAX_IMAGE_BBOX * NUM_BOX_ELEMENT).to_gpu();
vector<Job> fetch_jobs;
while(get_jobs_and_wait(fetch_jobs, max_batch_size)){
int infer_batch_size = fetch_jobs.size();
input->resize_single_dim(0, infer_batch_size);
for(int ibatch = 0; ibatch < infer_batch_size; ++ibatch){
auto& job = fetch_jobs[ibatch];
auto& mono = job.mono_tensor->data();
affin_matrix_device.copy_from_gpu(affin_matrix_device.offset(ibatch), mono->get_workspace()->gpu(), 6);
input->copy_from_gpu(input->offset(ibatch), mono->gpu(), mono->count());
job.mono_tensor->release();
}
engine->forward(false);
output_array_device.to_gpu(false);
for(int ibatch = 0; ibatch < infer_batch_size; ++ibatch){
auto& job = fetch_jobs[ibatch];
float* image_based_output = output->gpu<float>(ibatch);
float* output_array_ptr = output_array_device.gpu<float>(ibatch);
auto affine_matrix = affin_matrix_device.gpu<float>(ibatch);
checkCudaRuntime(cudaMemsetAsync(output_array_ptr, 0, sizeof(int), stream_));
decode_kernel_invoker(image_based_output, output->size(1), num_classes, confidence_threshold_, nms_threshold_, affine_matrix, output_array_ptr, MAX_IMAGE_BBOX, stream_);
}
output_array_device.to_cpu();
for(int ibatch = 0; ibatch < infer_batch_size; ++ibatch){
float* parray = output_array_device.cpu<float>(ibatch);
int count = min(MAX_IMAGE_BBOX, (int)*parray);
auto& job = fetch_jobs[ibatch];
auto& image_based_boxes = job.output;
for(int i = 0; i < count; ++i){
float* pbox = parray + 1 + i * NUM_BOX_ELEMENT;
int label = pbox[5];
int keepflag = pbox[6];
if(keepflag == 1){
image_based_boxes.emplace_back(pbox[0], pbox[1], pbox[2], pbox[3], pbox[4], label);
}
}
job.pro->set_value(image_based_boxes);
}
fetch_jobs.clear();
}
stream_ = nullptr;
tensor_allocator_.reset();
INFO("Engine destroy.");
}
virtual bool preprocess(Job& job, const Mat& image) override{
if(tensor_allocator_ == nullptr){
INFOE("tensor_allocator_ is nullptr");
return false;
}
job.mono_tensor = tensor_allocator_->query();
if(job.mono_tensor == nullptr){
INFOE("Tensor allocator query failed.");
return false;
}
AutoDevice auto_device(gpu_);
auto& tensor = job.mono_tensor->data();
if(tensor == nullptr){
// not init
tensor = make_shared<Tensor>();
tensor->set_workspace(make_shared<MixMemory>());
}
Size input_size(input_width_, input_height_);
job.additional.compute(image.size(), input_size);
tensor->set_stream(stream_);
tensor->resize(1, 3, input_height_, input_width_);
size_t size_image = image.cols * image.rows * 3;
size_t size_matrix = upbound(sizeof(job.additional.d2i), 32);
auto workspace = tensor->get_workspace();
uint8_t* gpu_workspace = (uint8_t*)workspace->gpu(size_matrix + size_image);
float* affine_matrix_device = (float*)gpu_workspace;
uint8_t* image_device = size_matrix + gpu_workspace;
uint8_t* cpu_workspace = (uint8_t*)workspace->cpu(size_matrix + size_image);
float* affine_matrix_host = (float*)cpu_workspace;
uint8_t* image_host = size_matrix + cpu_workspace;
//checkCudaRuntime(cudaMemcpyAsync(image_host, image.data, size_image, cudaMemcpyHostToHost, stream_));
// speed up
memcpy(image_host, image.data, size_image);
memcpy(affine_matrix_host, job.additional.d2i, sizeof(job.additional.d2i));
checkCudaRuntime(cudaMemcpyAsync(image_device, image_host, size_image, cudaMemcpyHostToDevice, stream_));
checkCudaRuntime(cudaMemcpyAsync(affine_matrix_device, affine_matrix_host, sizeof(job.additional.d2i), cudaMemcpyHostToDevice, stream_));
warp_affine_bilinear_and_normalize_plane(
image_device, image.cols * 3, image.cols, image.rows,
tensor->gpu<float>(), input_width_, input_height_,
affine_matrix_device, 114,
normalize_, stream_
);
return true;
}
virtual vector<shared_future<BoxArray>> commits(const vector<Mat>& images) override{
return ThreadSafedAsyncInferImpl::commits(images);
}
virtual std::shared_future<BoxArray> commit(const Mat& image) override{
return ThreadSafedAsyncInferImpl::commit(image);
}
private:
int input_width_ = 0;
int input_height_ = 0;
int gpu_ = 0;
float confidence_threshold_ = 0;
float nms_threshold_ = 0;
cudaStream_t stream_ = nullptr;
Norm normalize_;
};
void image_to_tensor(const cv::Mat& image, shared_ptr<Tensor>& tensor, Type type, int ibatch){
Norm normalize;
if(type == Type::V5){
normalize = Norm::alpha_beta(1 / 255.0f, 0.0f, ChannelType::SwapRB);
}else if(type == Type::X){
//float mean[] = {0.485, 0.456, 0.406};
//float std[] = {0.229, 0.224, 0.225};
//normalize_ = CUDAKernel::Norm::mean_std(mean, std, 1/255.0f, CUDAKernel::ChannelType::Invert);
normalize = Norm::None();
}else{
INFOE("Unsupport type %d", type);
}
Size input_size(tensor->size(3), tensor->size(2));
AffineMatrix affine;
affine.compute(image.size(), input_size);
size_t size_image = image.cols * image.rows * 3;
size_t size_matrix = upbound(sizeof(affine.d2i), 32);
auto workspace = tensor->get_workspace();
uint8_t* gpu_workspace = (uint8_t*)workspace->gpu(size_matrix + size_image);
float* affine_matrix_device = (float*)gpu_workspace;
uint8_t* image_device = size_matrix + gpu_workspace;
uint8_t* cpu_workspace = (uint8_t*)workspace->cpu(size_matrix + size_image);
float* affine_matrix_host = (float*)cpu_workspace;
uint8_t* image_host = size_matrix + cpu_workspace;
auto stream = tensor->get_stream();
memcpy(image_host, image.data, size_image);
memcpy(affine_matrix_host, affine.d2i, sizeof(affine.d2i));
checkCudaRuntime(cudaMemcpyAsync(image_device, image_host, size_image, cudaMemcpyHostToDevice, stream));
checkCudaRuntime(cudaMemcpyAsync(affine_matrix_device, affine_matrix_host, sizeof(affine.d2i), cudaMemcpyHostToDevice, stream));
warp_affine_bilinear_and_normalize_plane(
image_device, image.cols * 3, image.cols, image.rows,
tensor->gpu<float>(ibatch), input_size.width, input_size.height,
affine_matrix_device, 114,
normalize, stream
);
}
shared_ptr<Infer> create_infer(const string& engine_file, Type type, int gpuid, float confidence_threshold, float nms_threshold){
shared_ptr<YoloTRTInferImpl> instance(new YoloTRTInferImpl());
if(!instance->startup(engine_file, type, gpuid, confidence_threshold, nms_threshold)){
instance.reset();
}
return instance;
}
//Compile Model/
const char* mode_string(Mode type) {
switch (type) {
case Mode::FP32:
return "FP32";
case Mode::FP16:
return "FP16";
case Mode::INT8:
return "INT8";
default:
return "UnknowCompileMode";
}
}
typedef std::function<void(int current, int count, const std::vector<std::string>& files, std::shared_ptr<Tensor>& tensor)> Int8Process;
class Int8EntropyCalibrator : public IInt8EntropyCalibrator2{
public:
Int8EntropyCalibrator(const vector<string>& imagefiles, nvinfer1::Dims dims, const Int8Process& preprocess) {
Assert(preprocess != nullptr);
this->dims_ = dims;
this->allimgs_ = imagefiles;
this->preprocess_ = preprocess;
this->fromCalibratorData_ = false;
files_.resize(dims.d[0]);
checkCudaRuntime(cudaStreamCreate(&stream_));
}
Int8EntropyCalibrator(const vector<uint8_t>& entropyCalibratorData, nvinfer1::Dims dims, const Int8Process& preprocess) {
Assert(preprocess != nullptr);
this->dims_ = dims;
this->entropyCalibratorData_ = entropyCalibratorData;
this->preprocess_ = preprocess;
this->fromCalibratorData_ = true;
files_.resize(dims.d[0]);
checkCudaRuntime(cudaStreamCreate(&stream_));
}
virtual ~Int8EntropyCalibrator(){
checkCudaRuntime(cudaStreamDestroy(stream_));
}
int getBatchSize() const noexcept {
return dims_.d[0];
}
bool next() {
int batch_size = dims_.d[0];
if (cursor_ + batch_size > allimgs_.size())
return false;
int old_cursor = cursor_;
for(int i = 0; i < batch_size; ++i)
files_[i] = allimgs_[cursor_++];
if (!tensor_){
tensor_.reset(new Tensor(dims_.nbDims, dims_.d));
tensor_->set_stream(stream_);
tensor_->set_workspace(make_shared<MixMemory>());
}
preprocess_(old_cursor, allimgs_.size(), files_, tensor_);
return true;
}
bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept {
if (!next()) return false;
bindings[0] = tensor_->gpu();
return true;
}
const vector<uint8_t>& getEntropyCalibratorData() {
return entropyCalibratorData_;
}
const void* readCalibrationCache(size_t& length) noexcept {
if (fromCalibratorData_) {
length = this->entropyCalibratorData_.size();
return this->entropyCalibratorData_.data();
}
length = 0;
return nullptr;
}
virtual void writeCalibrationCache(const void* cache, size_t length) noexcept {
entropyCalibratorData_.assign((uint8_t*)cache, (uint8_t*)cache + length);
}
private:
Int8Process preprocess_;
vector<string> allimgs_;
size_t batchCudaSize_ = 0;
int cursor_ = 0;
nvinfer1::Dims dims_;
vector<string> files_;
shared_ptr<Tensor> tensor_;
vector<uint8_t> entropyCalibratorData_;
bool fromCalibratorData_ = false;
cudaStream_t stream_ = nullptr;
};
bool compile(
Mode mode, Type type,
unsigned int max_batch_size,
const string& source_onnx,
const string& saveto,
size_t max_workspace_size,
const std::string& int8_images_folder,
const std::string& int8_entropy_calibrator_cache_file) {
bool hasEntropyCalibrator = false;
vector<uint8_t> entropyCalibratorData;
vector<string> entropyCalibratorFiles;
auto int8process = [=](int current, int count, const vector<string>& files, shared_ptr<Tensor>& tensor){
for(int i = 0; i < files.size(); ++i){
auto& file = files[i];
INFO("Int8 load %d / %d, %s", current + i + 1, count, file.c_str());
auto image = cv::imread(file);
if(image.empty()){
INFOE("Load image failed, %s", file.c_str());
continue;
}
image_to_tensor(image, tensor, type, i);
}
tensor->synchronize();
};
if (mode == Mode::INT8) {
if (!int8_entropy_calibrator_cache_file.empty()) {
if (exists(int8_entropy_calibrator_cache_file)) {
entropyCalibratorData = load_file(int8_entropy_calibrator_cache_file);
if (entropyCalibratorData.empty()) {
INFOE("entropyCalibratorFile is set as: %s, but we read is empty.", int8_entropy_calibrator_cache_file.c_str());
return false;
}
hasEntropyCalibrator = true;
}
}
if (hasEntropyCalibrator) {
if (!int8_images_folder.empty()) {
INFOW("int8_images_folder is ignore, when int8_entropy_calibrator_cache_file is set");
}
}
else {
entropyCalibratorFiles = glob_image_files(int8_images_folder);
if (entropyCalibratorFiles.empty()) {
INFOE("Can not find any images(jpg/png/bmp/jpeg/tiff) from directory: %s", int8_images_folder.c_str());
return false;
}
if(entropyCalibratorFiles.size() < max_batch_size){
INFOW("Too few images provided, %d[provided] < %d[max batch size], image copy will be performed", entropyCalibratorFiles.size(), max_batch_size);
int old_size = entropyCalibratorFiles.size();
for(int i = old_size; i < max_batch_size; ++i)
entropyCalibratorFiles.push_back(entropyCalibratorFiles[i % old_size]);
}
}
}
else {
if (hasEntropyCalibrator) {
INFOW("int8_entropy_calibrator_cache_file is ignore, when Mode is '%s'", mode_string(mode));
}
}
INFO("Compile %s %s.", mode_string(mode), source_onnx.c_str());
shared_ptr<IBuilder> builder(createInferBuilder(gLogger), destroy_nvidia_pointer<IBuilder>);
if (builder == nullptr) {
INFOE("Can not create builder.");
return false;
}
shared_ptr<IBuilderConfig> config(builder->createBuilderConfig(), destroy_nvidia_pointer<IBuilderConfig>);
if (mode == Mode::FP16) {
if (!builder->platformHasFastFp16()) {
INFOW("Platform not have fast fp16 support");
}
config->setFlag(BuilderFlag::kFP16);
}
else if (mode == Mode::INT8) {
if (!builder->platformHasFastInt8()) {
INFOW("Platform not have fast int8 support");
}
config->setFlag(BuilderFlag::kINT8);
}
shared_ptr<INetworkDefinition> network;
shared_ptr<nvonnxparser::IParser> onnxParser;
const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
network = shared_ptr<INetworkDefinition>(builder->createNetworkV2(explicitBatch), destroy_nvidia_pointer<INetworkDefinition>);
//from onnx is not markOutput
onnxParser.reset(nvonnxparser::createParser(*network, gLogger), destroy_nvidia_pointer<nvonnxparser::IParser>);
if (onnxParser == nullptr) {
INFOE("Can not create parser.");
return false;
}
if (!onnxParser->parseFromFile(source_onnx.c_str(), 1)) {
INFOE("Can not parse OnnX file: %s", source_onnx.c_str());
return false;
}
auto inputTensor = network->getInput(0);
auto inputDims = inputTensor->getDimensions();
shared_ptr<Int8EntropyCalibrator> int8Calibrator;
if (mode == Mode::INT8) {
auto calibratorDims = inputDims;
calibratorDims.d[0] = max_batch_size;
if (hasEntropyCalibrator) {
INFO("Using exist entropy calibrator data[%d bytes]: %s", entropyCalibratorData.size(), int8_entropy_calibrator_cache_file.c_str());
int8Calibrator.reset(new Int8EntropyCalibrator(
entropyCalibratorData, calibratorDims, int8process
));
}
else {
INFO("Using image list[%d files]: %s", entropyCalibratorFiles.size(), int8_images_folder.c_str());
int8Calibrator.reset(new Int8EntropyCalibrator(
entropyCalibratorFiles, calibratorDims, int8process
));
}
config->setInt8Calibrator(int8Calibrator.get());
}
INFO("Input shape is %s", join_dims(vector<int>(inputDims.d, inputDims.d + inputDims.nbDims)).c_str());
INFO("Set max batch size = %d", max_batch_size);
INFO("Set max workspace size = %.2f MB", max_workspace_size / 1024.0f / 1024.0f);
int net_num_input = network->getNbInputs();
INFO("Network has %d inputs:", net_num_input);
vector<string> input_names(net_num_input);
for(int i = 0; i < net_num_input; ++i){
auto tensor = network->getInput(i);
auto dims = tensor->getDimensions();
auto dims_str = join_dims(vector<int>(dims.d, dims.d+dims.nbDims));
INFO(" %d.[%s] shape is %s", i, tensor->getName(), dims_str.c_str());
input_names[i] = tensor->getName();
}
int net_num_output = network->getNbOutputs();
INFO("Network has %d outputs:", net_num_output);
for(int i = 0; i < net_num_output; ++i){
auto tensor = network->getOutput(i);
auto dims = tensor->getDimensions();
auto dims_str = join_dims(vector<int>(dims.d, dims.d+dims.nbDims));
INFO(" %d.[%s] shape is %s", i, tensor->getName(), dims_str.c_str());
}
int net_num_layers = network->getNbLayers();
INFO("Network has %d layers", net_num_layers);
builder->setMaxBatchSize(max_batch_size);
config->setMaxWorkspaceSize(max_workspace_size);
auto profile = builder->createOptimizationProfile();
for(int i = 0; i < net_num_input; ++i){
auto input = network->getInput(i);
auto input_dims = input->getDimensions();
input_dims.d[0] = 1;
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, input_dims);
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, input_dims);
input_dims.d[0] = max_batch_size;
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, input_dims);
}
config->addOptimizationProfile(profile);
INFO("Building engine...");
auto time_start = chrono::duration_cast<chrono::milliseconds>(chrono::system_clock::now().time_since_epoch()).count();
shared_ptr<ICudaEngine> engine(builder->buildEngineWithConfig(*network, *config), destroy_nvidia_pointer<ICudaEngine>);
if (engine == nullptr) {
INFOE("engine is nullptr");
return false;
}
if (mode == Mode::INT8) {
if (!hasEntropyCalibrator) {
if (!int8_entropy_calibrator_cache_file.empty()) {
INFO("Save calibrator to: %s", int8_entropy_calibrator_cache_file.c_str());
save_file(int8_entropy_calibrator_cache_file, int8Calibrator->getEntropyCalibratorData());
}
else {
INFO("No set entropyCalibratorFile, and entropyCalibrator will not save.");
}
}
}
auto time_end = chrono::duration_cast<chrono::milliseconds>(chrono::system_clock::now().time_since_epoch()).count();
INFO("Build done %lld ms !", time_end - time_start);
// serialize the engine, then close everything down
shared_ptr<IHostMemory> seridata(engine->serialize(), destroy_nvidia_pointer<IHostMemory>);
return save_file(saveto, seridata->data(), seridata->size());
}
};
这段代码是 simple_yolo.hpp 头文件的实现部分，它包含了一些必要的库和头文件，以及一些辅助函数的实现。
其中，包含了 NvInfer.h、NvOnnxParser.h 和 cuda_runtime.h 头文件，它们是使用 NVIDIA TensorRT 进行推理所必需的库和头文件。
还包含了一些标准库和头文件，如 、、、 等，用于处理文件、内存等。
此外，还包含一些多线程相关的头文件，如 、
最后，还包含了一些条件编译指令，用于在不同的操作系统和编译器环境下进行编译和链接，以确保代码的可移植性和兼容性。
main.cpp
#include "simple_yolo.hpp"
#if defined(_WIN32)
# include 这是使用TensorRT引擎实现的YOLO目标检测算法的C++实现。该实现基于SimpleYolo库，该库提供了一个简单易用的接口来使用YOLO算法。
代码首先定义了一组实用函数，用于颜色转换、文件处理和性能测量。然后定义了一个名为requires_model()的函数，如果本地目录中不存在YOLO模型，则该函数下载模型。主函数main()调用这个函数下载模型，然后调用测试test()函数进行推理并测量性能。
测试函数test()设置设备ID和模式，然后调用SimpleYolo::compile()将ONNX模型编译为TensorRT引擎，如果引擎文件不存在于本地目录中。然后调用inference_and_performance()函数执行推理并测量模型的性能。结果以带有边界框和标签的图像形式保存在本地目录中。
总体来说，这个实现提供了一个简单而高效的方式，在C++中使用YOLO算法进行目标检测任务。