目录

1.CNN分类器

 1将分类器导出为ONNX格式

2.cpp推理

2.1：build_model()

2.2inference:

1.CNN分类器

学习实现一个完整的CNN分类器案例:

• 模型里面没有softmax操作，在这里采用包裹一层加上softmax节点后再导出模型，这样使得后处理得到的直接就是概率值，避免后处理上再做softmax
• 在c++代码中，则充分采用指针偏移的方式，提升cpu上预处理的效率
• 对于bgr->rgb也避免使用cvtColor实现，而是简单的改变赋值时的索引，提升效率

 1将分类器导出为ONNX格式

1.主干网络：将softmax分类直接加在主干网络resnet18后面

class Classifier(torch.nn.Module):def __init__(self):super().__init__()#使用torchvision自带的与训练模型, 更多模型请参考:https://tensorvision.readthedocs.io/en/master/self.backbone = torchvision.models.resnet18(pretrained=True)  def forward(self, x):feature     = self.backbone(x)probability = torch.softmax(feature, dim=1)return probability

2.图像预处理:

# 对每个通道进行归一化有助于模型的训练
imagenet_mean = [0.485, 0.456, 0.406]
imagenet_std  = [0.229, 0.224, 0.225]image = cv2.imread("workspace/dog.jpg")
image = cv2.resize(image, (224, 224))            # resize
image = image[..., ::-1]                         # BGR -> RGB
image = image / 255.0
image = (image - imagenet_mean) / imagenet_std   # normalize
image = image.astype(np.float32)                 # float64 -> float32
image = image.transpose(2, 0, 1)                 # HWC -> CHW
image = np.ascontiguousarray(image)              # contiguous array memory
image = image[None, ...]                         # CHW -> 1CHW
image = torch.from_numpy(image)                  # numpy -> torch

3.网络推理：model转为eval()模式

将图片输入网络得到概率和置信度

model = Classifier().eval()with torch.no_grad():probability   = model(image)predict_class = probability.argmax(dim=1).item()
confidence    = probability[0, predict_class]labels = open("workspace/labels.imagenet.txt").readlines()
labels = [item.strip() for item in labels]print(f"Predict: {predict_class}, {confidence}, {labels[predict_class]}")

4.导出onnx：

动态维度都设为：batch

opset_version=1

dummy = torch.zeros(1, 3, 224, 224)
torch.onnx.export(model, (dummy,), "workspace/classifier.onnx", input_names=["image"], output_names=["prob"], dynamic_axes={"image": {0: "batch"}, "prob": {0: "batch"}},opset_version=11
)

2.cpp推理

2.1 build_model()

1.智能指针：shared_prt管理内存的释放

// 通过智能指针管理nv返回的指针参数
// 内存自动释放，避免泄漏
template
shared_ptr<_T> make_nvshared(_T* ptr){return shared_ptr<_T>(ptr, [](_T* p){p->destroy();});
}
//这里用lambda 表达式的形式来表示 destroy 的方式

2.

bool build_model(){if(exists("engine.trtmodel")){printf("Engine.trtmodel has exists.\n");return true;}TRTLogger logger;// 这是基本需要的组件auto builder = make_nvshared(nvinfer1::createInferBuilder(logger));auto config = make_nvshared(builder->createBuilderConfig());auto network = make_nvshared(builder->createNetworkV2(1));// 通过onnxparser解析器解析的结果会填充到network中，类似addConv的方式添加进去auto parser = make_nvshared(nvonnxparser::createParser(*network, logger));if(!parser->parseFromFile("classifier.onnx", 1)){printf("Failed to parse classifier.onnx\n");// 注意这里的几个指针还没有释放，是有内存泄漏的，后面考虑更优雅的解决return false;}int maxBatchSize = 10;printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f);config->setMaxWorkspaceSize(1 << 28);// 如果模型有多个输入，则必须多个profileauto profile = builder->createOptimizationProfile();auto input_tensor = network->getInput(0);auto input_dims = input_tensor->getDimensions();// 配置最小、最优、最大范围input_dims.d[0] = 1;profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMIN, input_dims);profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kOPT, input_dims);input_dims.d[0] = maxBatchSize;profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMAX, input_dims);config->addOptimizationProfile(profile);auto engine = make_nvshared(builder->buildEngineWithConfig(*network, *config));if(engine == nullptr){printf("Build engine failed.\n");return false;}// 将模型序列化，并储存为文件auto model_data = make_nvshared(engine->serialize());FILE* f = fopen("engine.trtmodel", "wb");fwrite(model_data->data(), 1, model_data->size(), f);fclose(f);// 卸载顺序按照构建顺序倒序printf("Done.\n");return true;
}

2.2 inference:

1.构建engine等

  TRTLogger logger;auto engine_data = load_file("engine.trtmodel");auto runtime   = make_nvshared(nvinfer1::createInferRuntime(logger));auto engine = make_nvshared(runtime->deserializeCudaEngine(engine_data.data(), engine_data.size()));if(engine == nullptr){printf("Deserialize cuda engine failed.\n");runtime->destroy();return;}cudaStream_t stream = nullptr;checkRuntime(cudaStreamCreate(&stream));auto execution_context = make_nvshared(engine->createExecutionContext());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;float* input_data_host = nullptr;float* input_data_device = nullptr;checkRuntime(cudaMallocHost(&input_data_host, input_numel * sizeof(float)));checkRuntime(cudaMalloc(&input_data_device, input_numel * sizeof(float)));

2.加载图像,并预处理（bgr转为RGB用索引方式，比cvtColor快）

// image to floatauto image = cv::imread("dog.jpg");float mean[] = {0.406, 0.456, 0.485};float std[]  = {0.225, 0.224, 0.229};// 对应于pytorch的代码部分cv::resize(image, image, cv::Size(input_width, input_height));int image_area = image.cols * image.rows; //图像的面积unsigned char* pimage = image.data;  float* phost_b = input_data_host + image_area * 0;float* phost_g = input_data_host + image_area * 1;float* phost_r = input_data_host + image_area * 2; for(int i = 0; i < image_area; ++i, pimage += 3){// 注意这里的顺序rgb调换了*phost_r++ = (pimage[0] / 255.0f - mean[0]) / std[0];*phost_g++ = (pimage[1] / 255.0f - mean[1]) / std[1];*phost_b++ = (pimage[2] / 255.0f - mean[2]) / std[2];}

3.

    checkRuntime(cudaMemcpyAsync(input_data_device, input_data_host, input_numel * sizeof(float), cudaMemcpyHostToDevice, stream));// 3x3输入，对应3x3输出const int num_classes = 1000;float output_data_host[num_classes];float* output_data_device = nullptr;checkRuntime(cudaMalloc(&output_data_device, sizeof(output_data_host)));// 明确当前推理时，使用的数据输入大小auto input_dims = execution_context->getBindingDimensions(0);input_dims.d[0] = input_batch;// 设置当前推理时，input大小execution_context->setBindingDimensions(0, input_dims);float* bindings[] = {input_data_device, output_data_device};bool success      = execution_context->enqueueV2((void**)bindings, stream, nullptr);checkRuntime(cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream));checkRuntime(cudaStreamSynchronize(stream));float* prob = output_data_host;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);checkRuntime(cudaStreamDestroy(stream));checkRuntime(cudaFreeHost(input_data_host));checkRuntime(cudaFree(input_data_device));checkRuntime(cudaFree(output_data_device));

需要处理多个图像推理时：

    1. 在编译时，指定maxbatchsize为多个图

    2. 在推理时，指定输入的bindings shape的batch维度为使用的图像数，要求小于等于maxbatchsize

    3. 在收取结果的时候，tensor的shape是input指定的batch大小，按照batch处理即可

使用cudaMallocHost对输入的host进行分配，使得主机内存复制到设备效率更高：