關於深度學習實時檢測的三種方法(三)——Opencv DNN加載ONNX模型

碼道人 2021-08-15 20:42:03 阅读数:578

本文一共[544]字,预计阅读时长:1分钟~
深度 方法 opencv dnn onnx

簡介

隨著研究的逐漸深入,我了解到pytorch訓練模型再用libtorch加載實際上是一種很繞彎的方法,它的確能够滿足我們的需求,但也許不是最好的一種方法。如今比較通用的模型格式是ONNX模型,所以我就在想能否使用這種通用的模型來解决問題呢?

遠在天邊,近在眼前,沒想到OpenCV自己就有DNN模塊能够直接加載,我們要做的只需要將pytorch訓練出的pth模型轉成ONNX模型即可。本文章將介紹如何實現將pytorch訓練處的模型轉成ONNX模型並用DNN加載預測。

pth模型轉ONNX

上一篇文章我們介紹了如何用pytorch構建網絡並且訓練模型,訓練出的模型格式為pth,不清楚的朋友可以自行去看看上一篇文章,這裏接著上一篇,介紹一下如何轉模型的格式。

1.加載pth模型:

model = LeNet5(6)
state_dict = torch.load(input_pth_model, map_location='cpu').state_dict()
# load the model
model.load_state_dict(state_dict)

這裏說明一下,通過load().state_dict()加載的僅為模型的參數,能够更加輕量化節省內存提昇運行速度。如果是直接load則是直接加載整個模型的所有數據,所占用內存會較大一些。

2.轉化模型

dummy_input = torch.randn(1, 1, input_img_size, input_img_size)
input_names = ["input_image"]
output_names = ["output_classification"]
model.eval()
# 通過這裏轉化成onnx模型
torch.onnx.export(model, dummy_input, output_ONNX, verbose=True, input_names=input_names,
output_names=output_names)

其中input_img_size,input_pth_model,output_ONNX都需要自己指定。

加載ONNX模型

這裏我們首先創建一個類,能够達到管理網絡,加載模型並預測等功能。

類的創建

#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <iostream>
using namespace std;
using namespace cv;
class Dnn_NumDetect
{
public:
Dnn_NumDetect(const string& path);
// 運行前向傳遞以計算圖層的輸出
Point2f forward(Mat& src);
private:
dnn::Net Lenet5;
// 加載onnx模型
void loadModel(const string& path);
// 矩陣歸一化
void Mat_Normalization(Mat &matrix);
};

類的實現

#include "Dnn.h"
/** * @brief 使用Opencv Dnn Module 讀取ONNX模型 * @note 如果OpenCV是使用Intel的推理引擎庫編譯的,則DNN_BACKEND_DEFAULT錶示DNN_BACKEND_INFERENCE_ENGINE。 * 否則錶示DNN_BACKEND_OPENCV。 */
Dnn_NumDetect::Dnn_NumDetect(const string &path)
{
this->loadModel(path);
//網絡在支持的地方使用特定的計算後端
this->Lenet5.setPreferableBackend(dnn::DNN_BACKEND_OPENCV);
//網絡在特定的目標設備上進行計算
this->Lenet5.setPreferableTarget(dnn::DNN_TARGET_CPU);
}
/** * @brief 加載ONNX模型 */
void Dnn_NumDetect::loadModel(const string &path)
{
this->Lenet5 = dnn::readNetFromONNX(path);
CV_Assert(!this->Lenet5.empty());
}
/** * @brief 運行前向傳遞以計算圖層的輸出 * @return 指定層的第一個輸出的Blob。 */
Point2f Dnn_NumDetect::forward(Mat &src)
{
CV_Assert(!this->Lenet5.empty());
// 設置輸入
Mat input;
input = dnn::blobFromImage(src);
this->Lenet5.setInput(input);
Mat prob = this->Lenet5.forward();
// cout << prob <<endl;
// 矩陣歸一化
this->Mat_Normalization(prob);
// cout << prob <<endl;
Point classIdPoint;
double confidence;
//查找最大值和最小值
minMaxLoc(prob.reshape(1, 1), 0, &confidence, 0, &classIdPoint);
int classId = classIdPoint.x;
return Point2f(classId, confidence);
}

如果你需要用到更多關於DNN模塊的功能,可以直接去opencv官網的tutorial上查找,其中的dnn module中有更多詳細的介紹。

總結

關於數字識別,嘗試了三種方式:

1、用torch直接構建LeNet5網絡,產生pt模型

2、用pytorch訓練出pth模型後轉換成pt模型,在torch加載

3、用Opencv的dnn module 加載onnx模型(更方便用Openvino進行加速)

在計算運行時間的時候發現一個奇怪的現象,第一次預測時通常要花較長的時間,而之後運行時間將大幅度下降。

三種方式對比情况:

torch直接構建網絡訓練的時間最長,且准確率最低,不建議用這種方式。

torch加載pt模型和dnn module加載onnx模型准確率大致相同,但在運行時間上onnx模型更勝一籌。

參考內容

Deep Neural Networks (dnn module)

版权声明:本文为[碼道人]所创,转载请带上原文链接,感谢。 https://gsmany.com/2021/08/20210815204156364n.html