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.py的     cfg和     weights文件的路徑以及ONNX模型要保存的路徑。
• 執行     yolov3_to_onnx.py腳本，獲得     yolov3-tiny.onnx模型。

我們來看一下yolov3-tiny.onnx模型的可視化結果（使用Neutron），這里只看關鍵部分：

 

 

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

4. ONNX2TRT

在獲得了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。

 

5. 前向推理&后處理

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