前往小程序,Get更优阅读体验!
立即前往
首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >深度学习算法优化系列二十一 | 在VS2015上利用TensorRT部署YOLOV3-Tiny模型

深度学习算法优化系列二十一 | 在VS2015上利用TensorRT部署YOLOV3-Tiny模型

作者头像
BBuf
发布2020-04-01 11:06:24
1.2K0
发布2020-04-01 11:06:24
举报
文章被收录于专栏:GiantPandaCV

1. 前言

大家好,最近在VS2015上尝试用TensorRT来部署检测模型,中间走了两天弯路,感觉对于一个完全新手来说要做成功这件事并不会那么顺利。所以这里写一篇部署文章,希望能让使用TensorRT来部署YOLOV3-Tiny检测模型的同学少走一点弯路。

2. 确定走哪条路?

这里我是将AlexeyAB版本DarkNet训练出来的YOLOV3-Tiny检测模型(包含*.weights*.cfg)利用TensorRT部署在NVIDIA的1060显卡上。我选择的模型转换道路是DarkNet->ONNX->TRT。我们知道TensorRT既可以直接加载ONNX也可以加载ONNX转换得到的TRT引擎文件,而ONNX模型转TRT引擎文件是非常简单的,这个可以直接在代码里面完成,所以我们首先需要关注的是DarkNet模型转换到ONNX模型。

3. DarkNet2ONNX

现在已经明确,首要任务是模型转换为ONNX模型。这个我们借助Github上的一个工程就可以完成了,工程地址为:https://github.com/zombie0117/yolov3-tiny-onnx-TensorRT。具体操作步骤如下:

  • 克隆工程。
  • 使用Python2.7。
  • 执行pip install onnx=1.4.1
  • 将YOLOV3-Tiny的cfg文件末尾手动添加一个空行。
  • 修改yolov3_to_onnx.pycfgweights文件的路径以及ONNX模型要保存的路径。
  • 执行yolov3_to_onnx.py脚本,获得yolov3-tiny.onnx模型。

我们来看一下yolov3-tiny.onnx模型的可视化结果(使用Neutron),这里只看关键部分:

yolov3-tiny.onnx可视化

yolov3-tiny.onnx可视化

可以看到ONNX模型里面最后一层YOLO层是不存在了(ONNX不支持YOLO层,所以就忽略了),最后的两个输出层是做特征映射的卷积层,这就意味着后面的BBox和后处理NMS都是需要我们在代码中手动完成的。

4. ONNX2TRT

在获得了YOLOV3-Tiny的ONNX模型后,我们可以就可以将ONNX转为TensorRT的引擎文件了,这一转换的代码如下:

代码语言:javascript
复制
// 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。

yolov3-tiny.trt

5. 前向推理&后处理

这部分就没有什么细致讲解的必要了,直接给出源码吧。由于篇幅原因,我把源码上传到Github了。地址为:https://github.com/BBuf/cv_tools/blob/master/trt_yolov3_tiny.cpp注意我是用的TensorRT版本为6.0。修改ONNX模型的路径和图片路径就可以正确得到推理结果了。

6. 后记

这篇文章就是为大家分享了一个DarkNet2ONNX的工具,以及提供了在VS2015中利用TensorRT完成ONNX到TRT引擎文件的转换并进行预测的C++代码,希望可以帮助到刚使用TensorRT的同学,仅此而已。

本文参与 腾讯云自媒体同步曝光计划,分享自微信公众号。
原始发表:2020-03-28,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 GiantPandaCV 微信公众号,前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体同步曝光计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
目录
  • 1. 前言
  • 2. 确定走哪条路?
  • 3. DarkNet2ONNX
  • 4. ONNX2TRT
  • 5. 前向推理&后处理
  • 6. 后记
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档