使用OpenVINO加速Pytorch表情识别模型

小o 更新于 2年前

关于模型

模型是我基于前面一篇文章训练得到的,模型基于残差网络结构全卷积分类网络,
输入格式:NCHW=1x3x64x64
输出格式:NCHW=1x8x1x1

支持八种表情识别,列表如下:
[“neutral”,“anger”,“disdain”,“disgust”,“fear”,“happy”,“sadness”,“surprise”]

转ONNX

训练好的Pytorch模型可以保存为pt文件,通过pytorch自带的脚本可以转换为ONNX模型,这一步的转换脚本如下:

dummy_input = torch.randn(1, 3, 64, 64, device='cuda')
model = torch.load("./face_emotion***odel.pt")
output = model(dummy_input)
model.eval()
model.cuda()
torch.onnx.export(model, dummy_input, "face_emotion***odel.onnx", output_names={"output"}, verbose=True)

OpenCV DNN调用ONNX模型测试
转换为ONNX格式的模型,是可以通过OpenCV DNN模块直接调用的,调用方式如下:

emotion_net = cv.dnn.readNetFromONNX("face_emotion***odel.onnx")
image = cv.imread("D:/facedb/test/367.jpg")
blob = cv.dnn.blobFromImage(image, 0.00392, (64, 64), (0.5, 0.5, 0.5), False) / 0.5
emotion_net.setInput(blob)
res = emotion_net.forward("output")
idx = np.argmax(np.reshape(res, (8)))
emotion_txt = emotion_labels[idx]
cv.putText(image, emotion_txt, (10, 25), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
cv.imshow("input", image)
cv.waitKey(0)
cv.destroyAllWindow******lockquote>
运行结果如下:


ONNX转IR

如何把ONNX文件转换OpenVINO的IR文件?答案是借助OpenVINO的模型优化器组件工具,OpenVINO的模型优化器组件工具支持常见的Pytorch预训练模型与torchvision迁移训练模型的转换,

要转换ONNX到IR,首先需要安装ONNX组件支持,直接运行OpenVINO预安装脚本即可获得支持,截图如下:

然后执行下面的转换脚本即可:

不用怀疑了,转换成功!

使用OpenVINO的Inference Engine加速推理
对得到的模型通过OpenVINO安装包自带的OpenCV DNN完成调用,设置加速推理引擎为Inference Engine,这部分的代码如下:

dnn::Net emtion_net = readNetFromModelOptimizer(emotion_xml, emotion_bin);
emtion_net.setPreferableTarget(DNN_TARGET_CPU);
emtion_net.setPreferableBackend(DNN_BACKEND_INFERENCE_ENGINE);

执行推理与输出解析,得到表情分类的结果,代码如下:

Rect box(x1, y1, x2 - x1, y2 - y1);
Mat roi = frame(box);
Mat face_blob = blobFromImage(roi, 0.00392, Size(64, 64), Scalar(0.5, 0.5, 0.5), false, false);
emtion_net.setInput(face_blob);
Mat probs = emtion_net.forward();
int index = 0;
float max = -1;
for (int i = 0; i < 8; i++) {
const float *scores = probs.ptr<float>(0, i, 0);
float score = scores[0];
if (max < score) {
max = score;
index = i;
}
}

最终的运行结果如下图:


0个评论