如何在VS2015上利用TensorRT部署YOLOV3-Tiny模型

1. 前言

大家好，最近在VS2015上嘗試用TensorRT來(lái)部署檢測(cè)模型，中間走了兩天彎路，感覺(jué)對(duì)于一個(gè)完全新手來(lái)說(shuō)要做成功這件事并不會(huì)那么順利。所以這里寫一篇部署文章，希望能讓使用TensorRT來(lái)部署YOLOV3-Tiny檢測(cè)模型的同學(xué)少走一點(diǎn)彎路。

2. 確定走哪條路？

這里我是將AlexeyAB版本DarkNet訓(xùn)練出來(lái)的YOLOV3-Tiny檢測(cè)模型（包含*.weights和*.cfg）利用TensorRT部署在NVIDIA的1060顯卡上。我選擇的模型轉(zhuǎn)換道路是DarkNet->ONNX->TRT。我們知道TensorRT既可以直接加載ONNX也可以加載ONNX轉(zhuǎn)換得到的TRT引擎文件，而ONNX模型轉(zhuǎn)TRT引擎文件是非常簡(jiǎn)單的，這個(gè)可以直接在代碼里面完成，所以我們首先需要關(guān)注的是DarkNet模型轉(zhuǎn)換到ONNX模型。

3. DarkNet2ONNX

現(xiàn)在已經(jīng)明確，首要任務(wù)是模型轉(zhuǎn)換為ONNX模型。這個(gè)我們借助Github上的一個(gè)工程就可以完成了，工程地址為：https://github.com/zombie0117/yolov3-tiny-onnx-TensorRT。具體操作步驟如下：

• 克隆工程。
• 使用Python2.7。
• 執(zhí)行     pip install onnx=1.4.1
• 將YOLOV3-Tiny的     cfg文件末尾手動(dòng)添加一個(gè)空行。
• 修改     yolov3_to_onnx.py的     cfg和     weights文件的路徑以及ONNX模型要保存的路徑。
• 執(zhí)行     yolov3_to_onnx.py腳本，獲得     yolov3-tiny.onnx模型。

我們來(lái)看一下yolov3-tiny.onnx模型的可視化結(jié)果（使用Neutron），這里只看關(guān)鍵部分：

 

 

可以看到ONNX模型里面最后一層YOLO層是不存在了(ONNX不支持YOLO層，所以就忽略了)，最后的兩個(gè)輸出層是做特征映射的    卷積層，這就意味著后面的BBox和后處理NMS都是需要我們?cè)诖a中手動(dòng)完成的。

4. ONNX2TRT

在獲得了YOLOV3-Tiny的ONNX模型后，我們可以就可以將ONNX轉(zhuǎn)為TensorRT的引擎文件了，這一轉(zhuǎn)換的代碼如下：

// ONNX模型轉(zhuǎn)為TensorRT引擎
bool onnxToTRTModel(const std::string& modelFile, // onnx文件的名字
	const std::string& filename,  // TensorRT引擎的名字
	IHostMemory*& trtModelStream) // output buffer for the TensorRT model
{
	// 創(chuàng)建builder
	IBuilder* builder = createInferBuilder(gLogger.getTRTLogger());
	assert(builder != nullptr);
	nvinfer1::INetworkDefinition* network = builder->createNetwork();

	// 解析ONNX模型
	auto parser = nvonnxparser::createParser(*network, gLogger.getTRTLogger());


	//可選的 - 取消下面的注釋可以查看網(wǎng)絡(luò)中每層的星系信息
	//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();

	// 銷毀所有相關(guān)的東西
	engine->destroy();
	network->destroy();
	builder->destroy();

	return true;
}

執(zhí)行了這個(gè)函數(shù)之后就會(huì)在指定的目錄下生成yolov3-tiny.trt，從下圖可以看到這個(gè)引擎文件有48.6M，而原始的weights文件是34.3M。

 

5. 前向推理&后處理

這部分就沒(méi)有什么細(xì)致講解的必要了，直接給出源碼吧。由于篇幅原因，我把源碼上傳到Github了。地址為：https://github.com/BBuf/cv_tools/blob/master/trt_yolov3_tiny.cpp。注意我是用的TensorRT版本為6.0。修改ONNX模型的路徑和圖片路徑就可以正確得到推理結(jié)果了。