今天就跟大家聊聊有关如何在VS2015上利用TensorRT部署YOLOV3-Tiny模型,可能很多人都不太了解,为了让大家更加了解,小编给大家总结了以下内容,希望大家根据这篇文章可以有所收获。
大家好,最近在VS2015上尝试用TensorRT来部署检测模型,中间走了两天弯路,感觉对于一个完全新手来说要做成功这件事并不会那么顺利。所以这里写一篇部署文章,希望能让使用TensorRT来部署YOLOV3-Tiny检测模型的同学少走一点弯路。
这里我是将AlexeyAB版本DarkNet训练出来的YOLOV3-Tiny检测模型(包含*.weights
和*.cfg
)利用TensorRT部署在NVIDIA的1060显卡上。我选择的模型转换道路是DarkNet->ONNX->TRT
。我们知道TensorRT既可以直接加载ONNX也可以加载ONNX转换得到的TRT引擎文件,而ONNX模型转TRT引擎文件是非常简单的,这个可以直接在代码里面完成,所以我们首先需要关注的是DarkNet模型转换到ONNX模型。
现在已经明确,首要任务是模型转换为ONNX模型。这个我们借助Github上的一个工程就可以完成了,工程地址为:https://github.com/zombie0117/yolov3-tiny-onnx-TensorRT
。具体操作步骤如下:
pip install onnx=1.4.1
cfg
文件末尾手动添加一个空行。yolov3_to_onnx.py
的
cfg
和
weights
文件的路径以及ONNX模型要保存的路径。yolov3_to_onnx.py
脚本,获得
yolov3-tiny.onnx
模型。我们来看一下yolov3-tiny.onnx
模型的可视化结果(使用Neutron),这里只看关键部分:
可以看到ONNX模型里面最后一层YOLO
层是不存在了(ONNX不支持YOLO
层,所以就忽略了),最后的两个输出层是做特征映射的
卷积层,这就意味着后面的BBox和后处理NMS都是需要我们在代码中手动完成的。
在获得了YOLOV3-Tiny的ONNX模型后,我们可以就可以将ONNX转为TensorRT的引擎文件了,这一转换的代码如下:
// ONNX模型转为TensorRT引擎bool onnxToTRTModel(const std::string& modelFile, // onnx文件的名字 const std::string& filename, // TensorRT引擎的名字 IHostMemory*& trtModelStream) // output buffer for the TensorRT model{ // 创建builder IBuilder* builder = createInferBuilder(gLogger.getTRTLogger()); assert(builder != nullptr); nvinfer1::INetworkDefinition* network = builder->createNetwork(); // 解析ONNX模型 auto parser = nvonnxparser::createParser(*network, gLogger.getTRTLogger()); //可选的 - 取消下面的注释可以查看网络中每层的星系信息 //config->setPrintLayerInfo(true); //parser->reportParsingInfo(); //判断是否成功解析ONNX模型 if (!parser->parseFromFile(modelFile.c_str(), static_cast<int>(gLogger.getReportableSeverity()))) { gLogError << "Failure while parsing ONNX file" << std::endl; return false; } // 建立推理引擎 builder->setMaxBatchSize(BATCH_SIZE); builder->setMaxWorkspaceSize(1 << 30); builder->setFp16Mode(true); builder->setInt8Mode(gArgs.runInInt8); if (gArgs.runInInt8) { samplesCommon::setAllTensorScales(network, 127.0f, 127.0f); } cout << "start building engine" << endl; ICudaEngine* engine = builder->buildCudaEngine(*network); cout << "build engine done" << endl; assert(engine); // 销毁模型解释器 parser->destroy(); // 序列化引擎 trtModelStream = engine->serialize(); // 保存引擎 nvinfer1::IHostMemory* data = engine->serialize(); std::ofstream file; file.open(filename, std::ios::binary | std::ios::out); cout << "writing engine file..." << endl; file.write((const char*)data->data(), data->size()); cout << "save engine file done" << endl; file.close(); // 销毁所有相关的东西 engine->destroy(); network->destroy(); builder->destroy(); return true;}
执行了这个函数之后就会在指定的目录下生成yolov3-tiny.trt
,从下图可以看到这个引擎文件有48.6M,而原始的weights
文件是34.3M。
这部分就没有什么细致讲解的必要了,直接给出源码吧。由于篇幅原因,我把源码上传到Github了。地址为:https://github.com/BBuf/cv_tools/blob/master/trt_yolov3_tiny.cpp
。注意我是用的TensorRT版本为6.0。修改ONNX模型的路径和图片路径就可以正确得到推理结果了。
看完上述内容,你们对如何在VS2015上利用TensorRT部署YOLOV3-Tiny模型有进一步的了解吗?如果还想了解更多知识或者相关内容,请关注亿速云行业资讯频道,感谢大家的支持。
亿速云「云服务器」,即开即用、新一代英特尔至强铂金CPU、三副本存储NVMe SSD云盘,价格低至29元/月。点击查看>>
免责声明:本站发布的内容(图片、视频和文字)以原创、转载和分享为主,文章观点不代表本网站立场,如果涉及侵权请联系站长邮箱:is@yisu.com进行举报,并提供相关证据,一经查实,将立刻删除涉嫌侵权内容。
原文链接:https://my.oschina.net/u/4580321/blog/4358400