我想用 OpenCV-DNN 包装注意力 OCR模型以增加推理时间。我正在使用官方 TF 模型repo中的 TF 代码。
对于使用 OpenCV-DNN 包装 TF 模型,我指的是此代码。需要“cv2.dnn.readNetFromTensorflow()
冻结图”和“图结构”来读取 TF 模型。
我使用此代码片段从元检查点文件导入结构并将图形结构保存在.pbtxt
文件中。
# load graph from meta file
tf.reset_default_graph()
imported_meta = tf.train.import_meta_graph("attention_ocr_2017_08_09/model_demo_inference.ckpt.meta")
# restore graph structure, variables in session's graph
sess = tf.Session()
imported_meta.restore(sess, 'attention_ocr_2017_08_09/model_demo_inference.ckpt')
# write graph structure to a pbtxt file
tf.train.write_graph(sess.graph_def, './', 'train_attention.pbtxt', as_text=True)
冻结图,代码如下:
from tensorflow.python.tools import freeze_graph
freeze_graph.freeze_graph('train_attention.pbtxt', '', False, \
'attention_ocr_2017_08_09/model_demo_inference.ckpt', \
'AttentionOcr_v1_1/Softmax', \
'save/restore_all', 'save/Const:0', 'frozen_model.pb', True, "")
最终代码使用函数中的pbtxt
和pb
文件cv2.dnn.readNetFromTensorflow()
。
# Wrap TF model in OpenCV DNN
import cv2
FROZEN_GRAPH = "frozen_model.pb"
PB_TXT = "train_attention.pbtxt"
img = cv2.imread('testdata/fsns_train_00.png')
blob = cv2.dnn.blobFromImage(img,1)
net = cv2.dnn.readNetFromTensorflow(FROZEN_GRAPH, PB_TXT)
out = net.forward()
out
遇到的错误是:
---------------------------------------------------------------------------
error Traceback (most recent call last)
<ipython-input-128-09e46e8b88ed> in <module>
9 blob = cv2.dnn.blobFromImage(img,1)
10
---> 11 net = cv2.dnn.readNetFromTensorflow(FROZEN_GRAPH, PB_TXT)
12 out = net.forward()
13 out
error: OpenCV(4.0.0) /Users/travis/build/skvark/opencv-python/opencv/modules/dnn/src/
tensorflow/tf_io.cpp:54: error: (-2:Unspecified error)
FAILED: ReadProtoFromTextFile(param_file, param).
Failed to parse GraphDef file: train_attention.pbtxt in function 'ReadTFNetParamsFromTextFileOrDie'
注意:输出节点名称是通过查看生成的图中的张量列表手动设置的:
# get names of all tensors
def get_names(graph=sess.graph):
return [t.name for op in graph.get_operations() for t in op.values()]
l1 = get_names()
for ele in l1:
print(ele)
我将非常感谢 SO 社区提供的任何帮助。