TRT4-trt-integrate - 4 TRT的封装
前言：
- 对tensorRT的封装，更像是对推理引擎的封装
- 封装的意义在于对技术的标准化、工具化，能够是的使用时更加便利，效率更高，定制更多的默认行为
- 封装推理引擎的思想，还可以应用到更多其他地方。嵌入式、等等。由于大多推理引擎提供的默认方式不够友好，对其进行包装，能够很好的使得自己的代码具有复用性，一套代码多处用
- 还可以实现，同样的封装，通过简单的配置，切换不同的推理后端。这都取决于需求
- 我们的唯一目的就是让工作更简单，让代码复用性更强，让技术可以沉淀
builder的封装：
头文件：
namespace TRT { enum class Mode : int { FP32, FP16 }; const char* mode_string(Mode type); bool compile( Mode mode, unsigned int maxBatchSize, const std::string& source, const std::string& saveto, const size_t maxWorkspaceSize = 1ul << 30 // 1ul << 30 = 1GB ); };这里主要就是输入5个参数，野花似我们build model最重要的5个参数。
分别是:
mode ：决定使用什么精度
maxBatchsize ：最大的batch是多少
source ： 要推理的onnx文件名称
saveto ： 要保存的trtmodel名称
maxWorkspaceSize : 运行工作空间
compile函数：
bool compile(
Mode mode,
unsigned int maxBatchSize,
const string& source,
const string& saveto,
const size_t maxWorkspaceSize) {
INFO("Compile %s %s.", mode_string(mode), source.c_str());
auto builder = make_nvshared(createInferBuilder(gLogger));
if (builder == nullptr) {
INFOE("Can not create builder.");
return false;
}
auto config = make_nvshared(builder->createBuilderConfig());
if (mode == Mode::FP16) {
if (!builder->platformHasFastFp16()) {
INFOW("Platform not have fast fp16 support");
}
config->setFlag(BuilderFlag::kFP16);
}
shared_ptr network;
//shared_ptr caffeParser;
const auto explicitBatch = 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
network = make_nvshared(builder->createNetworkV2(explicitBatch));
shared_ptr onnxParser = make_nvshared(nvonnxparser::createParser(*network, gLogger));
if (onnxParser == nullptr) {
INFOE("Can not create parser.");
return false;
}
if (!onnxParser->parseFromFile(source.c_str(), 1)) {
INFOE("Can not parse OnnX file: %s", source.c_str());
return false;
}
auto inputTensor = network->getInput(0);
auto inputDims = inputTensor->getDimensions();
INFO("Input shape is %s", join_dims(vector(inputDims.d, inputDims.d + inputDims.nbDims)).c_str());
INFO("Set max batch size = %d", maxBatchSize);
INFO("Set max workspace size = %.2f MB", maxWorkspaceSize / 1024.0f / 1024.0f);
INFO("Base device: %s", CUDATools::device_description().c_str());
int net_num_input = network->getNbInputs();
INFO("Network has %d inputs:", net_num_input);
vector 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(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(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);
//打印详细信息
for(int i = 0; i < net_num_layers; ++i){
auto layer = network->getLayer(i);
auto name = layer->getName();
auto type_str = layer_type_name(layer);
auto input0 = layer->getInput(0);
if(input0 == nullptr) continue;
auto output0 = layer->getOutput(0);
auto input_dims = input0->getDimensions();
auto output_dims = output0->getDimensions();
bool has_input = layer_has_input_tensor(layer);
bool has_output = layer_has_output_tensor(layer);
auto descript = layer_descript(layer);
type_str = align_blank(type_str, 18);
auto input_dims_str = align_blank(dims_str(input_dims), 18);
auto output_dims_str = align_blank(dims_str(output_dims), 18);
auto number_str = align_blank(format("%d.", i), 4);
const char* token = " ";
if(has_input)
token = " >>> ";
else if(has_output)
token = " *** ";
INFOV("%s%s%s %s-> %s%s", token,
number_str.c_str(),
type_str.c_str(),
input_dims_str.c_str(),
output_dims_str.c_str(),
descript.c_str()
);
}
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(maxWorkspaceSize);
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] = maxBatchSize;
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, input_dims);
}
// not need
// for(int i = 0; i < net_num_output; ++i){
// auto output = network->getOutput(i);
// auto output_dims = output->getDimensions();
// output_dims.d[0] = 1;
// profile->setDimensions(output->getName(), nvinfer1::OptProfileSelector::kMIN, output_dims);
// profile->setDimensions(output->getName(), nvinfer1::OptProfileSelector::kOPT, output_dims);
// output_dims.d[0] = maxBatchSize;
// profile->setDimensions(output->getName(), nvinfer1::OptProfileSelector::kMAX, output_dims);
// }
config->addOptimizationProfile(profile);
// error on jetson
// auto timing_cache = shared_ptr(config->createTimingCache(nullptr, 0), [](nvinfer1::ITimingCache* ptr){ptr->reset();});
// config->setTimingCache(*timing_cache, false);
// config->setFlag(BuilderFlag::kGPU_FALLBACK);
// config->setDefaultDeviceType(DeviceType::kDLA);
// config->setDLACore(0);
INFO("Building engine...");
auto time_start = timestamp_now();
auto engine = make_nvshared(builder->buildEngineWithConfig(*network, *config));
if (engine == nullptr) {
INFOE("engine is nullptr");
return false;
}
INFO("Build done %lld ms !", timestamp_now() - time_start);
// serialize the engine, then close everything down
auto seridata = make_nvshared(engine->serialize());
return save_file(saveto, seridata->data(), seridata->size());
}
}; //namespace TRTBuilder
看上去很庞大，但其实中间一大块：
for(int i = 0; i < net_num_layers; ++i){
auto layer = network->getLayer(i);
auto name = layer->getName();
auto type_str = layer_type_name(layer);
auto input0 = layer->getInput(0);
if(input0 == nullptr) continue;
auto output0 = layer->getOutput(0);
auto input_dims = input0->getDimensions();
auto output_dims = output0->getDimensions();
bool has_input = layer_has_input_tensor(layer);
bool has_output = layer_has_output_tensor(layer);
auto descript = layer_descript(layer);
type_str = align_blank(type_str, 18);
auto input_dims_str = align_blank(dims_str(input_dims), 18);
auto output_dims_str = align_blank(dims_str(output_dims), 18);
auto number_str = align_blank(format("%d.", i), 4);
const char* token = " ";
if(has_input)
token = " >>> ";
else if(has_output)
token = " *** ";
INFOV("%s%s%s %s-> %s%s", token,
number_str.c_str(),
type_str.c_str(),
input_dims_str.c_str(),
output_dims_str.c_str(),
descript.c_str()
);
}
都是在打印详细信息。
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(maxWorkspaceSize);
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] = maxBatchSize;
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, input_dims);
}
// not need
// for(int i = 0; i < net_num_output; ++i){
// auto output = network->getOutput(i);
// auto output_dims = output->getDimensions();
// output_dims.d[0] = 1;
// profile->setDimensions(output->getName(), nvinfer1::OptProfileSelector::kMIN, output_dims);
// profile->setDimensions(output->getName(), nvinfer1::OptProfileSelector::kOPT, output_dims);
// output_dims.d[0] = maxBatchSize;
// profile->setDimensions(output->getName(), nvinfer1::OptProfileSelector::kMAX, output_dims);
// }
config->addOptimizationProfile(profile);
// error on jetson
// auto timing_cache = shared_ptr(config->createTimingCache(nullptr, 0), [](nvinfer1::ITimingCache* ptr){ptr->reset();});
// config->setTimingCache(*timing_cache, false);
// config->setFlag(BuilderFlag::kGPU_FALLBACK);
// config->setDefaultDeviceType(DeviceType::kDLA);
// config->setDLACore(0);
INFO("Building engine...");
auto time_start = timestamp_now();
auto engine = make_nvshared(builder->buildEngineWithConfig(*network, *config));
if (engine == nullptr) {
INFOE("engine is nullptr");
return false;
}
INFO("Build done %lld ms !", timestamp_now() - time_start);
// serialize the engine, then close everything down
auto seridata = make_nvshared(engine->serialize());
return save_file(saveto, seridata->data(), seridata->size());
} 真正的这一块和咱们之前的没什么太大的区别。
main函数调用
bool build_model(){
if(exists("engine.trtmodel")){
printf("Engine.trtmodel has exists.\n");
return true;
}
//SimpleLogger::set_log_level(SimpleLogger::LogLevel::Verbose);//打印详细信息
TRT::compile(
TRT::Mode::FP32,
10,
"classifier.onnx",
"engine.trtmodel",
1 << 28
);
INFO("Done.");
return true;
}比之前省了很多的事情。
Memory的封装
通常我们在定义input_device的时候，还会定义一块host上的内存，而这两者的大小近乎是完全一样的，所以既然着两者是成对出现的，我们将其打包到一起会不会更方便呢？
头文件：
class MixMemory {
public:
MixMemory(int device_id = CURRENT_DEVICE_ID);
MixMemory(void* cpu, size_t cpu_size, void* gpu, size_t gpu_size, int device_id = CURRENT_DEVICE_ID);
virtual ~MixMemory();
void* gpu(size_t size);
void* cpu(size_t size);
template
_T* gpu(size_t size){ return (_T*)gpu(size * sizeof(_T)); }
template
_T* cpu(size_t size){ return (_T*)cpu(size * sizeof(_T)); };
void release_gpu();
void release_cpu();
void release_all();
// 是否属于我自己分配的gpu/cpu 下面的函数几乎都只是为了方便，可以不用看
inline bool owner_gpu() const{return owner_gpu_;}
inline bool owner_cpu() const{return owner_cpu_;}
inline size_t cpu_size() const{return cpu_size_;}
inline size_t gpu_size() const{return gpu_size_;}
inline int device_id() const{return device_id_;}
inline void* gpu() const { return gpu_; }
// Pinned Memory
inline void* cpu() const { return cpu_; }
template
inline _T* gpu() const { return (_T*)gpu_; }
// Pinned Memory
template
inline _T* cpu() const { return (_T*)cpu_; }
void reference_data(void* cpu, size_t cpu_size, void* gpu, size_t gpu_size, int device_id = CURRENT_DEVICE_ID);
private:
void* cpu_ = nullptr;
size_t cpu_size_ = 0;
bool owner_cpu_ = true;
int device_id_ = 0;
void* gpu_ = nullptr;
size_t gpu_size_ = 0;
bool owner_gpu_ = true;
}; GPU内存分配和复用：
这个原理其实特别简单：
分配函数：
void* MixMemory::gpu(size_t size) {
if (gpu_size_ < size) {
release_gpu();
gpu_size_ = size;
CUDATools::AutoDevice auto_device_exchange(device_id_);
checkRuntime(cudaMalloc(&gpu_, size));
checkRuntime(cudaMemset(gpu_, 0, size));
}
return gpu_;
}就是比大小，如果之前你要了100字节，这次只要了50字节，那我根本就不用重新分配，直接把上一次的给你就好了，反正也没释放。但是如果这一次要了1000字节，笔之前大了好多，那我就将之前的先释放掉，再重新分配给你。
释放函数：
void MixMemory::release_gpu() {
if (gpu_) {
if(owner_gpu_){
CUDATools::AutoDevice auto_device_exchange(device_id_);
checkRuntime(cudaFree(gpu_));
}
gpu_ = nullptr;
}
gpu_size_ = 0;
}CPU内存分配和复用
和GPU一样：
void* MixMemory::cpu(size_t size) {
if (cpu_size_ < size) {
release_cpu();
cpu_size_ = size;
CUDATools::AutoDevice auto_device_exchange(device_id_);
checkRuntime(cudaMallocHost(&cpu_, size));
assert(cpu_ != nullptr);
memset(cpu_, 0, size);
}
return cpu_;
}
void MixMemory::release_cpu() {
if (cpu_) {
if(owner_cpu_){
CUDATools::AutoDevice auto_device_exchange(device_id_);
checkRuntime(cudaFreeHost(cpu_));
}
cpu_ = nullptr;
}
cpu_size_ = 0;
}析构函数：
void MixMemory::release_all() {
release_cpu();
release_gpu();
}MixMemory::~MixMemory() {
release_all();
}
再退出作用域的时候自动释放掉。
Tensor的封装：
对tensor进行封装，张量是CNN中常见的基本单元，尤其是计算偏移量的工作需要封装，其次是内存的复制、分配需要引用memory进行包装，避免使用时面对指针不好管控
- 1 内存的管理，可以使用mixmemory进行解决
- 2 内存的复用，依然可以使用mixmemory解决
- 3 内存的拷贝，CPU->GPU or GPU->CPU
解决方案：
a 定义内存的状态，表明内存的最新内容在哪里（GPU/CPU/Init）
b 懒分配原则，不使用即使空壳，只有需要用到的时候再分配
c 获取内存地址，即表示想拿到最新的数据。 比如说tensor.cpu，就表示我想拿到最新的数据，并且放到CPU上。
- 4 索引的计算，比如一个5d的Tensor(B,D,C,H,W)。此时要获取B = 1 ， D = 3 ， C = 0 ，H = 5 ， W = 9的位置元素，就需要计算索引，这也是非常基本的一个能力。
定义内存状态：
enum class DataHead : int{
Init = 0,
Device = 1,
Host = 2
};
enum class DataType : int {
Float = 0,
Float16 = 1,
Int32 = 2,
UInt8 = 3
};内存的拷贝：
拿这个构造函数举个例子：
Tensor::Tensor(DataType dtype, shared_ptr data, int device_id){
shape_string_[0] = 0;
descriptor_string_[0] = 0;
this->device_id_ = get_device(device_id);
dtype_ = dtype;
setup_data(data);
} void Tensor::setup_data(shared_ptr data){
data_ = data;
if(data_ == nullptr){
data_ = make_shared(device_id_);
}else{
device_id_ = data_->device_id();
}
head_ = DataHead::Init;
if(data_->cpu()){
head_ = DataHead::Host;
}
if(data_->gpu()){
head_ = DataHead::Device;
}
} 这个setup_data表示，如果船进来为空，则重新分配一下，如果船进来的不为空，那么就将device_id直接给device_id_。
如果data->cpu不为空，那就代表host端已经有数据了，那么就可以将这个head定为Host，如果是在GPU上，那就定义为Device,无论如何，这个head永远指向最新的地址。
既然我们之前说过了，一旦我们想获取内存地址，就是想获取最新的数据，比如tensor.cpu那就代表要获得数据，并且放到CPU上。
inline void* cpu() const { ((Tensor*)this)->to_cpu(); return data_->cpu(); }那我们就来看一下to_cpu
Tensor& Tensor::to_cpu(bool copy) {
if (head_ == DataHead::Host)
return *this;
head_ = DataHead::Host;
data_->cpu(bytes_);
if (copy && data_->gpu() != nullptr) {
CUDATools::AutoDevice auto_device_exchange(this->device());
checkRuntime(cudaMemcpyAsync(data_->cpu(), data_->gpu(), bytes_, cudaMemcpyDeviceToHost, stream_));
checkRuntime(cudaStreamSynchronize(stream_));
}
return *this;
}
如果head == host ，那就表明现在就在CPU上，那么直接返回就可以。
但是如果不在的话，那就需要先对host内存作分配。
如果标记为需要拷贝，同时GPU是不为null，就是代表有数据（因为如果GPU是空的，那就代表没有数据，那既然没有数据，那自然也是不需要拷贝的）。
之后我们用cudaMemcpyAsync将GPU数据拷贝到CPU上。
拷贝结束，返回this。再对于异步拷贝集上流同步。
这样就完成了，我们在使用的时候根本不需要考虑这是什么内存，只要我在我想要的时候访问，我就能拿到想要的数据。
索引计算
shape/dim index
B 1
D 3
C 0
H 5
W 9
position = 0 for d,i in zip(dims , indexs): position *= d position += i左乘右加
int offset(int index, _Args ... index_args) const{
const int index_array[] = {index, index_args...};
return offset_array(sizeof...(index_args) + 1, index_array);
}这个函数参数是可变的，也就是一个变参，之后将这个参数塞给array。
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;
}传进来的参数就是index_array，之后对这个数组进行咱们上面的运算。算好之后就是我们的索引。
main函数调用
int input_batch = 1;
int input_channel = 3;
int input_height = 224;
int input_width = 224;
int input_numel = input_batch * input_channel * input_height * input_width;
// tensor的建立并不会立即分配内存，而是在第一次需要使用的时候进行分配
TRT::Tensor input_data({input_batch, input_channel, input_height, input_width}, TRT::DataType::Float);
// 为input关联stream，使得在同一个pipeline中执行复制操作
input_data.set_stream(stream);
可以直接通过shape进行初始化，之后setstream与stream关联，拷贝的时候就可以异步啦。
// 利用opencv mat的内存地址引用，实现input与mat的关联，然后利用split函数一次性完成mat到input的复制
cv::Mat channel_based[3];
for(int i = 0; i < 3; ++i)
// 注意这里 2 - i是实现bgr -> rgb的方式
// 这里cpu提供的参数0是表示batch的索引是0，第二个参数表示通道的索引，因此获取的是0, 2-i通道的地址
// 而tensor最大的好处就是帮忙计算索引，否则手动计算就得写很多代码
channel_based[i] = cv::Mat(input_height, input_width, CV_32F, input_data.cpu(0, 2-i));
cv::split(image, channel_based);//image就被送到channel_based里面去了。
// 利用opencv的mat操作加速减去均值和除以标准差
for(int i = 0; i < 3; ++i)
channel_based[i] = (channel_based[i] / 255.0f - mean[i]) / std[i];
float* bindings[] = {input_data.gpu(), output_data.gpu()};
bool success = execution_context->enqueueV2((void**)bindings, stream, nullptr);
checkRuntime(cudaStreamSynchronize(stream)); 推理的时候直接给到GPU
避免了这些冗余的操作：
cudaMalloc(&input_data_device, sizeof(input_data_host)); cudaMalloc(&output_data_device, sizeof(output_data_host)); cudaMemcpyAsync(input_data_device, input_data_host, sizeof(input_data_host), cudaMemcpyHostToDevice, stream); // 用一个指针数组指定input和output在gpu中的指针。 float* bindings[] = {input_data_device, output_data_device};
// 当获取cpu地址的时候，如果数据最新的在gpu上，就进行数据复制，然后再返回cpu地址
float* prob = output_data.cpu(); 推理完成后再拿回到CPU
同样避免了：
bool success = execution_context->enqueueV2((void**)bindings, stream, nullptr); cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream); cudaStreamSynchronize(stream);
Infer的封装：
对infer进行封装，有了基本组件，可以拼接一个完整的推理器，而且该推理器的思想可以应用到很多框架作为底层，并不只限制于tensorRT，还可以是rknn、openvino等。
void InferImpl::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();
}
}和之前一样的，先创建一个context上下文，getbinding，获取尺寸信息。之后运行enqueueV2进行推理，而后对其进行同步。
main函数调用：
engine->forward(true);
int num_classes = output->size(1);
float* prob = output->cpu();
int predict_label = std::max_element(prob, prob + num_classes) - prob;
auto labels = load_labels("labels.imagenet.txt");
auto predict_name = labels[predict_label];
float confidence = prob[predict_label];
printf("Predict: %s, confidence = %f, label = %d\n", predict_name.c_str(), confidence, predict_label); 在main函数里只有这么短短的一行：
engine->forward(true); 至此，TRT的封装结束。