我终于找到了解决我的问题的方法。想法是使用 Keras 声明模型，使用 TensorFlow API 将其转换为 protobuf 图定义，然后深入生成的图定义以删除不支持的节点。
我把所有东西都封装到了下面的python文件中：
import tensorflow as tf
from tensorflow.core import framework
def find_all_nodes(graph_def, **kwargs):
for node in graph_def.node:
for key, value in kwargs.items():
if getattr(node, key) != value:
break
else:
yield node
raise StopIteration
def find_node(graph_def, **kwargs):
try:
return next(find_all_nodes(graph_def, **kwargs))
except StopIteration:
raise ValueError(
'no node with attributes: {}'.format(
', '.join("'{}': {}".format(k, v) for k, v in kwargs.items())))
def walk_node_ancestors(graph_def, node_def, exclude=set()):
openlist = list(node_def.input)
closelist = set()
while openlist:
name = openlist.pop()
if name not in exclude:
node = find_node(graph_def, name=name)
openlist += list(node.input)
closelist.add(name)
return closelist
def remove_nodes_by_name(graph_def, node_names):
for i in reversed(range(len(graph_def.node))):
if graph_def.node[i].name in node_names:
del graph_def.node[i]
def make_shape_node_const(node_def, tensor_values):
node_def.op = 'Const'
node_def.ClearField('input')
node_def.attr.clear()
node_def.attr['dtype'].type = framework.types_pb2.DT_INT32
tensor = node_def.attr['value'].tensor
tensor.dtype = framework.types_pb2.DT_INT32
tensor.tensor_shape.dim.add()
tensor.tensor_shape.dim[0].size = len(tensor_values)
for value in tensor_values:
tensor.tensor_content += value.to_bytes(4, 'little')
output_shape = node_def.attr['_output_shapes']
output_shape.list.shape.add()
output_shape.list.shape[0].dim.add()
output_shape.list.shape[0].dim[0].size = len(tensor_values)
def make_cv2_compatible(graph_def):
# A reshape node needs a shape node as its second input to know how it
# should reshape its input tensor.
# When exporting a model using Keras, this shape node is computed
# dynamically using `Shape`, `StridedSlice` and `Pack` operators.
# Unfortunately those operators are not supported yet by the OpenCV API.
# The goal here is to remove all those unsupported nodes and hard-code the
# shape layer as a const tensor instead.
for reshape_node in find_all_nodes(graph_def, op='Reshape'):
# Get a reference to the shape node
shape_node = find_node(graph_def, name=reshape_node.input[1])
# Find and remove all unsupported nodes
garbage_nodes = walk_node_ancestors(graph_def, shape_node,
exclude=[reshape_node.input[0]])
remove_nodes_by_name(graph_def, garbage_nodes)
# Infer the shape tensor from the reshape output tensor shape
if not '_output_shapes' in reshape_node.attr:
raise AttributeError(
'cannot infer the shape node value from the reshape node. '
'Please set the `add_shapes` argument to `True` when calling '
'the `Session.graph.as_graph_def` method.')
output_shape = reshape_node.attr['_output_shapes'].list.shape[0]
output_shape = [dim.size for dim in output_shape.dim]
# Hard-code the inferred shape in the shape node
make_shape_node_const(shape_node, output_shape[1:])
有了这个，我可以执行以下脚本：
import tensorflow as tf
from tensorflow.python.keras.models import Sequential
from tensorflow.python.keras.layers import Dense, Activation, Conv2D, MaxPooling2D, Flatten, Reshape
from tensorflow.python.keras import backend as K
import numpy as np
import graph_util
# Define the model in Keras
model = Sequential()
model.add(Conv2D(32,kernel_size=(3,3),input_shape=(28,28,1)))
model.add(Activation('relu'))
model.add(Conv2D(32,kernel_size=(3,3)))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(2,2)))
model.add(Flatten())
model.add(Dense(128))
model.add(Activation('relu'))
model.add(Dense(10))
model.add(Activation('softmax'))
model.summary()
# Get Keras prediction
inp = np.random.standard_normal([1, 28, 28, 1]).astype(np.float32)
tf_out = model.predict(inp)
# Serialize and fix the graph
sess = K.get_session()
graph_def = sess.graph.as_graph_def(add_shapes=True)
graph_def = tf.graph_util.convert_variables_to_constants(sess, graph_def, [model.output.name.split(':')[0]])
graph_util.make_cv2_compatible(graph_def)
# Print the graph nodes
print('\n'.join(node.name for node in graph_def.node))
# Save the graph as a binary protobuf2 file
tf.train.write_graph(graph_def, '', 'model.pb', as_text=False)
# Get OpenCV prediction
import cv2 as cv
net = cv.dnn.readNetFromTensorflow('model.pb')
net.setInput(inp.transpose(0, 3, 1, 2))
cv_out = net.forward()
print(np.max(np.abs(tf_out - cv_out)))
输出在 1e-7、1e-8 范围内