【OpenVINO】将TensorFlow模型转成IR文件，并部署到NCS2上运行

TensorFlow模型

这里以本人用TensorFlow实现的AlexNet模型为例，详见文章用TensorFlow实现AlexNet。

且下面的过程都在激活OpenVINO环境后打开的jupyter notebook中完成。

模型转换

编写好AlexNet模型并通过测试后，将模型专为pb文件:

inputs = tf.placeholder(tf.float32, [1, 227, 227, 3])

model = AlexNetModel()
score = model.inference(inputs)
softmax = tf.nn.softmax(score, name='prob')

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    model.load_original_weights(sess)
    
    constant_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, ['prob'])
    
    with tf.gfile.GFile('frozen_alexnet_model.pb', mode='wb') as f:
        f.write(constant_graph.SerializeToString())

之后在命令行中，将pb文件用OpenVINO转换为IR文件，格式如下为：

python (mo_tf.py路径) --input_model (pb文件路径) --output_dir (IR文件输出路径) --data_type (FP32或FP16) --batch (输入批量大小）

例如本人在win10环境下生成适用于CPU/GPU的IR模型，执行的命令为：

python “C:\Program Files (x86)\IntelSWTools\openvino\deployment_tools\model_optimizer\mo_tf.py” --input_model frozen_alexnet_model.pb --output_dir lrmodels/FP32 --data_type FP32 --batch 1

生成适用于NCS的IR模型，执行的命令为:

python “C:\Program Files (x86)\IntelSWTools\openvino\deployment_tools\model_optimizer\mo_tf.py” --input_model frozen_alexnet_model.pb --output_dir lrmodels/FP16 --data_type FP16 --batch 1

更多通用的转换参数详见这里，TensorFlow模型特有的转换参数则详见这里。

使用模型

转换好的IR模型需要使用OpenVINO的python库来调用

from openvino.inference_engine import IENetwork, IEPlugin # 导入openvino库
from caffe_classes import class_names
import numpy as np
import matplotlib.pyplot as plt
import cv2
import os

imagenet_mean = np.array([104., 117., 124.], dtype=np.float32)

# 读入图片并进行一些处理
image = cv2.imread('llama.jpeg') 
img = cv2.resize(image.astype(np.float32), (227, 227))
img -= imagenet_mean
img = img.reshape((1, 227, 227, 3))
img = img.transpose((0, 3, 1, 2)) 

model_xml_CPU = 'lrmodels/FP32/frozen_alexnet_model.xml' # xml文件路径
model_bin_CPU = os.path.splitext(model_xml_CPU)[0] + ".bin" # bin文件路径

model_xml_MYRIAD = 'lrmodels/FP16/frozen_alexnet_model.xml'
model_bin_MYRIAD = os.path.splitext(model_xml_MYRIAD)[0] + ".bin"

plugin = IEPlugin(device='CPU') # 使用CPU运行模型，要用NCS的话改为MYRIAD
net = IENetwork(model=model_xml_CPU, weights=model_bin_CPU) # 用NCS的话也需要更改
net.batch_size = 1 # 批量
input_blob = next(iter(net.inputs)) # 迭代器
exec_net = plugin.load(network=net) # 载入模型
outputs = exec_net.infer(inputs={input_blob: img}) # 推理

class_name = class_names[np.argmax(outputs['prob'])]
probs = outputs['prob'][0, np.argmax(outputs['prob'])]

plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
plt.title("Class: " + class_name + ", probability: %.4f" %probs)

运行后得到的测试结果与原模型相同：
[外链图片转存失败(img-9QOSrNld-1563196834838)(https://i.loli.net/2019/05/06/5ccf87c57fdab.png)]
所以将训练好的模型用OpenVINO部署到NCS等设备上的流程差不多都是这样，更深入的内容还需要研读OpenVINO的官方文档。

参考资料

1. 完整程序-Github
2. OpenVINO官方文档