目录
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<typename _T> 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); // 如果模型有多个输入,则必须多个profile auto 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 float auto 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进行分配,使得主机内存复制到设备效率更高: