您好,登錄后才能下訂單哦!
今天就跟大家聊聊有關如何在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模型有進一步的了解嗎?如果還想了解更多知識或者相關內容,請關注億速云行業資訊頻道,感謝大家的支持。
免責聲明:本站發布的內容(圖片、視頻和文字)以原創、轉載和分享為主,文章觀點不代表本網站立場,如果涉及侵權請聯系站長郵箱:is@yisu.com進行舉報,并提供相關證據,一經查實,將立刻刪除涉嫌侵權內容。