细说TensorRT C++模型部署3

相关文章:

• 细说TensorRT C++部署1
• 细说TensorRT C++部署2
• 细说TensorRT C++部署4

五 线程的具体实现

对 (四) 中提到的三个线程的代码实现. 这是最重要的三个代码片段了.

• trtPre代码解析:

void InferImpl::trtPre(BaseParam &curParam, Infer *curFunc) {
    // 记录预处理总耗时
    double totalTime;
    std::chrono::system_clock::time_point preTime;
//    auto t = timer.curTimePoint();
//    计算1个batchSize的数据所需的空间大小
    unsigned long mallocSize = this->memory[0] * sizeof(float);
    // count统计推理图片数量,最后一次,可能小于batchSize.imgNums统计预处理图片数量,index是gpuIn的索引,两个显存轮换使用
    int countPre = 0, index = 0, inputSize = this->memory[0] / curParam.batchSize;
    Job job;

    while (workRunning) {
        {
            std::unique_lock<std::mutex> l1(lock1);
            // 队列不为空, 退出等待,. 当标志位workRunning为False时,说明要停止线程运行,也退出等待
            cv_.wait(l1, [&]() { return !qfJobs.empty() || !workRunning; });
//            判断退出等待原因是否是要停止线程
            if (!workRunning) break;
            fJob = qfJobs.front();
            qfJobs.pop();
        }
        // 默认推理的数量为batchSize, 只有最后一次才可能小于batchSize
        job.inferNum = curParam.batchSize;
//        记录待推理的最后一个元素地址
        auto lastElement = &fJob.mats.back();

        // todo 暂时,先在内存中进行图片预处理. gpuMat以后写cuda核函数处理
        for (auto &curMat: fJob.mats) {
            // 记录前处理耗时
            preTime = timer.curTimePoint();
//            调用具体算法自身实现的预处理逻辑
            curFunc->preProcess(curParam, curMat, pinMemoryIn + countPre * inputSize);
            // 将当前正在处理图片的变换参数加入该batch的变换vector中
            job.d2is.push_back({curParam.d2i[0], curParam.d2i[1], curParam.d2i[2], curParam.d2i[3], curParam.d2i[4], curParam.d2i[5]});

            countPre += 1;
            // 不是最后一个元素且数量小于batchSize,继续循环,向pinMemoryIn写入预处理后数据
            if (&curMat != lastElement && countPre < curParam.batchSize) continue;
            // 若是最后一个元素,记录当前需要推理图片数量(可能小于一个batchSize)
            if (&curMat == lastElement) job.inferNum = countPre;

            totalTime += timer.timeCountS(preTime);
            //若全部在gpu上操作,不需要这句复制
            checkRuntime(cudaMemcpyAsync(gpuIn[index], pinMemoryIn, mallocSize, cudaMemcpyHostToDevice, preStream));
            job.inputTensor = gpuIn[index];

            countPre = 0;
//            将内存地址索引指向另一块内存
            index = index >= 1 ? 0 : index + 1;

            {
                std::unique_lock<std::mutex> l1(lock1);
                // false,继续等待. true,不等待,跳出wait. 一旦进入wait,解锁. 退出wait,又加锁 最多2个batch
                cv_.wait(l1, [&]() { return qPreJobs.size() < 2; });
                // 将一个batch待推理数据的显存空间指针保存
                qPreJobs.push(job);
            }

            // 流同步,在通知队列可使用前,确保待推理数据已复制到gpu上,保证在推理时取出就能用
            checkRuntime(cudaStreamSynchronize(preStream));
            cv_.notify_all();
            // 清空保存仿射变换参数,只保留当前推理batch的参数
            std::vector<std::vector<float >>().swap(job.d2is);
        }
    }
    // 结束预处理线程,释放资源
    preFinish = true;
    // 唤醒trt线程,告知预处理线程已结束
    cv_.notify_all();
//    printf("pre   use time: %.2f ms\n", totalTime);
    logInfo("pre   use time: %.3f s", totalTime);
}

在while循环中,如果队列qfJobs不为空, 就提取数据, 之后遍历数据容器内每一项进行preProcess数据预处理.
这里的预处理是指由各具体算法自己实现的预处理, 将处理好的数据存放在内存pinMemoryIn开头的地址中,当处理完一个batchsize数据后, 再向后进行其他步骤. 若python中传入多个batchsize数据,且处理到最后一个数据且不满一个batchsize时, 代码也会向下进行. 此时job.inferNum会单独记录当前待推理数据数量.
数据处理完毕后,会使用cuda函数cudaMemcpyAsync把数据异步复制到GPU上, 然后传入队列qPreJobs, 当流同步后唤醒下一个线程trtInfer.
这里要说明, 在gpu上开辟的空间有两个空间,分别对应gpuIn[0],gpuIn[1],每个空间可以存储一个bachsize处理好的数据, 这样的处理方式比较占用显存, 好处是能使队列qPreJobs中一直有可用数据, 即使推理模块trtInfer速度远快于队列数据生成, 那么也能保证trtProcess一直处于工作状态, 保证效率.

• trtInfer代码解析:

// 适用于模型推理的通用trt流程
void InferImpl::trtInfer(BaseParam &curParam) {
    // 记录推理耗时
    double totalTime;
//    auto t = timer.curTimePoint();
    int index2 = 0;
//    engine输入输出节点名字, 是把model编译为onnx文件时,在onnx中手动写的输入输出节点名
    const char *inferInputName = curParam.inputName.c_str();
    const char *inferOutputName = curParam.outputName.c_str();

    Job job;
    Out outTrt;

    while (true) {
        {
            std::unique_lock<std::mutex> l1(lock1);
            // 队列不为空, 就说明推理空间图片已准备好,退出等待,继续推理. 当图片都处理完,并且队列为空,要退出等待,因为此时推理工作已完成
            cv_.wait(l1, [&]() { return !qPreJobs.empty() || (preFinish && qPreJobs.empty()); });
//            若图片都处理完且队列空,说明推理结束,直接退出线程
            if (preFinish && qPreJobs.empty()) break;
            job = qPreJobs.front();
            qPreJobs.pop();
        }
        // 消费掉一个元素,通知队列跳出等待,向qJob继续写入一个batch的待推理数据
        cv_.notify_all();

        auto qT1 = timer.curTimePoint();
        // 指定tensor的输入输出地址,然后进行推理
        curParam.context->setTensorAddress(inferInputName, job.inputTensor);
        curParam.context->setTensorAddress(inferOutputName, gpuOut[index2]);
        curParam.context->enqueueV3(inferStream);

        outTrt.inferOut = gpuOut[index2];
        outTrt.d2is = job.d2is;
        outTrt.inferNum = job.inferNum;

        index2 = index2 >= 1 ? 0 : index2 + 1;

        cudaStreamSynchronize(inferStream);
        totalTime += timer.timeCountS(qT1);

        // 流同步后,获取该batchSize推理结果
        {
            std::unique_lock<std::mutex> l2(lock2);
            // false, 表示继续等待. true, 表示不等待,gpuOut内只有两块空间,因此队列长度只能限定为2
            cv_.wait(l2, [&]() { return qPostJobs.size() < 2; });
            qPostJobs.emplace(outTrt);
        }
        cv_.notify_all();
    }
    // 在post后处理线程中判断,所有推理流程是否结束,然后决定是否结束后处理线程
    inferFinish = true;
    //告知post后处理线程,推理线程已结束
    cv_.notify_all();
//    printf("infer use time: %.2f ms\n", totalTime);
    logInfo("infer use time: %.3f s", totalTime);
}

这个函数思路明确, 从预处理队列取数据, 然后指定引擎输入输出节点的地址和数据值, 执行推理, 流同步后将推理结果传入后处理队列qPostJobs中.
该推理模块中也是设计了两个显存块存储推理结果, 存储在qPostJobs中,设计思路在trtPre中有说明.

• trtPost代码解析:

void InferImpl::trtPost(BaseParam &curParam, Infer *curFunc) {
    // 记录后处理耗时
    double totalTime;
//    auto t = timer.curTimePoint();
//    将推理后数据从从显存拷贝到内存中,计算所需内存大小,ps:其实与占用显存大小一致
    unsigned long mallocSize = this->memory[1] * sizeof(float), singleOutputSize = this->memory[1] / curParam.batchSize;
//    batchBox收集每个batchSize后处理结果,然后汇总到batchBoxes中
    batchBoxesType batchBoxes, batchBox;

//    传入图片总数, 已处理图片数量
    int totalNum = 0, countPost = 0;
    bool flag = true;
    Out outPost;

    while (true) {
        {
            std::unique_lock<std::mutex> l2(lock2);
            // 队列不为空, 就说明图片已推理好,退出等待,进行后处理. 推理结束,并且队列为空,退出等待,因为此时推理工作已完成
            cv_.wait(l2, [&]() { return !qPostJobs.empty() || (inferFinish && qPostJobs.empty()); });
            // 退出推理线程
            if (inferFinish && qPostJobs.empty()) break;
            outPost = qPostJobs.front();
            qPostJobs.pop();
            cv_.notify_all();
//           每次当python传入数据时,totalNum值才会取出一次,获得当前传入图片数量
            if (flag) {
                totalNum = qfJobLength.front();
                qfJobLength.pop();
                flag = false;
            }
        }
        // todo 将engine推理好的结果从gpu转移到内存中处理, 更高效的方式是在在gpu中用cuda核函数处理,but,以后再扩展吧
        cudaMemcpy(pinMemoryOut, outPost.inferOut, mallocSize, cudaMemcpyDeviceToHost);
        // 取回数据的仿射变换量,用于还原预测量在原图尺寸上的位置
        curParam.d2is = outPost.d2is;

        auto qT1 = timer.curTimePoint();
//        记录当前后处理图片数量, 若是单张图片,这个记录没啥用. 若是传入多个batchSize的图片,countPost会起到标识作用
        countPost += outPost.inferNum;
//        调用具体算法自身实现的后处理逻辑
        curFunc->postProcess(curParam, pinMemoryOut, singleOutputSize, outPost.inferNum, batchBox);
        //将每次后处理结果合并到输出vector中
        batchBoxes.insert(batchBoxes.end(), batchBox.begin(), batchBox.end());
        batchBox.clear();
        // 当commit中传入图片处理完时,通过set_value返回所有图片结果. 重新计数, 并返回下一次要输出推理结果的图片数量
        if (totalNum <= countPost) {
//            fJob.batchResult = std::make_shared<std::promise<batchBoxesType>>();
            // 输出解码后的结果,在commit中接收
            fJob.batchResult->set_value(batchBoxes);
            countPost = 0;
            flag = true;
            batchBoxes.clear();
        }
        totalTime += timer.timeCountS(qT1);
    }
//    printf("post  use time: %.2f ms\n", totalTime);
    logInfo("post  use time: %.3f s", totalTime);

}

代码中变量flag是标志位,作用时什么呢?
在commit方法中qfJobLength队列记录了传入数据的量,在trtPost中取出, 当后处理数据量达到totalNum时, 说明传入的数据全部后处理完毕,可以返回出去了.
flag初始为True,当第一次执行trtPost时取出传入数据的数量大小. 当后处理结束时再置为True,等待从python传入的新数据.
设想这样一个场景, python中通过request 请求一次传入66张图片,但batchsize仅设为8,这时需处理9次, 最后一次处理数量为2,当进行到后处理时, flag为True,会获得totalNum=66, 后处理会将每次处理的结果量记录到countPost中, 然后合并到vector变量batchBoxes中. 当totalNum <= countPost时, 说明66张图片全部后处理处理完,通过set_value()返回结果,然后清空计数,重置flag状态.