jetson nx tensorRT: 使用engine文件进行网络推理(语意分割,内存零copy,同步异步inference)

写3种方法,1.含CPU与GPU间的内存copy,异步inference,2. 内存零copy, 异步inference, 3. 内存零copy, 同步inference

  1. 含CPU与GPU间的内存copy,异步inference
    大多代码都用的是这种方式

先声明一个input和output数组,

static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W]; //input
static float prob[BATCH_SIZE * OUTPUT_SIZE];  //output

读入engine文件

char *trtModelStream{nullptr};
string engine_name = "net.engine";
ifstream file(engine_name, std::ios::binary);
if (file.good()) {
   file.seekg(0, file.end);
   size = file.tellg();
   file.seekg(0, file.beg);
   trtModelStream = new char[size];
   assert(trtModelStream);
   file.read(trtModelStream, size);
   file.close();
}

准备cudaEngine

IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;

读入图片

cv::Mat img = cv::imread(string(argv[2]) + "/" + file_names[count]);  //读一个文件夹下的图片,文件夹路径由参数传入

把img的数据传入data数组,这个要依具体需求而定,可能要预处理一下,这里略过
然后把input, output的指针传入inference函数

doInference(*context, data, prob, BATCH_SIZE);

doInference的实现细节如下:

void doInference(IExecutionContext& context, float* input, float* output, int batchSize)
{
    const ICudaEngine& engine = context.getEngine();

    // Pointers to input and output device buffers to pass to engine.
    // Engine requires exactly IEngine::getNbBindings() number of buffers.
    assert(engine.getNbBindings() == 2);
    void* buffers[2];

    // In order to bind the buffers, we need to know the names of the input and output tensors.
    // Note that indices are guaranteed to be less than IEngine::getNbBindings()
    const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
    const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);

    // Create GPU buffers on device
    CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
    CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));

    // Create stream
    cudaStream_t stream;
    CHECK(cudaStreamCreate(&stream));

    //这里把input的数据copy到buffer的input中,即把数据从CPU拷贝到GPU
    CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
    //inference
    context.enqueue(batchSize, buffers, stream, nullptr);
    //把推理的结果从buffer copy到output中,即把数据从GPU copy到CPU,这步耗时较多
    CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
   //做同步
    cudaStreamSynchronize(stream);
    
    // Release stream and buffers
    cudaStreamDestroy(stream);
    CHECK(cudaFree(buffers[inputIndex]));
    CHECK(cudaFree(buffers[outputIndex]));
}

可以看到这里要为buffer保留内存,然后把input的数据从CPU copy到buffer中,推理完之后再把buffer中的结果copy到output

output这个指针传入的是prob,那么可以直接利用prob得到网络输出的结果,做一个softmax, 即得到分割结果
下面再用阈值分割出前景和背景,并把结果保存出来

cv::Mat mask_mat = cv::Mat(INPUT_H, INPUT_W, CV_8UC1);
char* ptmp = NULL;
for(int i = 0; i < INPUT_H; i++) {
  ptmp = mask_mat.ptr<uchar>(i);
  for(int j = 0; j < INPUT_W; j++) {
     float numerator = exp(prob[i*INPUT_W + j]);
     float pixel = numerator / (numerator + exp(prob[i*INPUT_W + j + INPUT_W*INPUT_H]) + 0.0001); //softmax
     if(pixel > CONF_THRESH) {
          ptmp[j] = 255;
     } else {
          ptmp[j] = 0;
     }
   }
 }

cv::imwrite(filename, mask_mat);
  1. 内存零copy,异步inference
    由于nx板子的CPU和GPU内存是共享的,不需要在它们之间copy数据

主要流程和上面一样,区别在于input和output数组,buffer的声明
input和output仅声明为指针,并保留内存

float* data;
float* prob;
CHECK(cudaHostAlloc((void **)&data, BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float), cudaHostAllocMapped));
CHECK(cudaHostAlloc((void **)&prob, BATCH_SIZE * OUTPUT_SIZE * sizeof(float), cudaHostAllocMapped));

参数传入还是一样的

doInference_zeroCopy(*context, data, prob, BATCH_SIZE);

doInference_zeroCopy的实现如下

void doInference_zeroCopy(IExecutionContext& context, float* input, float* output, int batchSize)
{
    cudaSetDeviceFlags(cudaDeviceMapHost);
    const ICudaEngine& engine = context.getEngine();
    // Pointers to input and output device buffers to pass to engine.
    // Engine requires exactly IEngine::getNbBindings() number of buffers.
    assert(engine.getNbBindings() == 2);
    void* buffers[2];

    const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
    const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
	
		//buffer中仅传入指针
    cudaHostGetDevicePointer((void**)&buffers[inputIndex], (void*)input, 0);
    cudaHostGetDevicePointer((void**)&buffers[outputIndex], (void*)output, 0);
		
		// Create stream
    cudaStream_t stream;
    CHECK(cudaStreamCreate(&stream));
    
    context.enqueue(batchSize, buffers, stream, nullptr);
    
    cudaStreamSynchronize(stream);
    
    // Release stream and buffers
    cudaStreamDestroy(stream);
}
  1. 内存零copy,同步inference
    与异步inference的区别仅在于把doInference_zeroCopy函数中
context.enqueue(batchSize, buffers, stream, nullptr);

替换成

context.execute(batchSize, buffers);

由于不再需要stream,可以把stream的部分注释掉

经过测试,发现影响时间效率的主要在于图片的尺寸大小,
另外,在图片size 640x640时,同步推理比异步推理时间效率略高一些

评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

蓝羽飞鸟

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值