之前写过一篇文章有关tensorflow ckpt和pb模型之间转换的操作, 这次再详细说下里面的一些坑.
ckpt 转 pb
如果只有ckpt文件,例如三个ckpt文件model-xxx.data, model-xxx.meta, model-xxx.index
, 转成pb文件方式很简单,按照上面的链接就可以,只需要确定好输出节点名称就可以。如果不知道,可以打印ckpt里面的节点名称,或者使用tf.summary.FileWriter
保存模型用tensorboard去查看。
我在转一个模型的时候,发现打印ckpt里面的节点名称,有两种方法:
# read node name way 1
import tensorflow as tf
with tf.Session() as sess:
saver = tf.train.import_meta_graph(ckpt + '.meta', clear_devices=True)
graph_def = tf.get_default_graph().as_graph_def(add_shapes=True)
node_list=[n.name for n in graph_def.node]
for node in node_list:
print (node)
#read node name way 2
#the two ways result different
from tensorflow.python import pywrap_tensorflow
reader = pywrap_tensorflow.NewCheckpointReader(ckpt)
var_to_shape_map = reader.get_variable_to_shape_map()
for key in var_to_shape_map:
print("tensor_name: ", key)
但是这两种方法打印出来的节点名称不一样,按照经验来说,这两种方式应该打印的一样,我也没有找到原因在哪。还有一个方法是直接使用graph.get_operations()
,结果和第一种一致。
pb转tvm
因为生成pb后需要转成tvm,如果模型里面有resize layer
,那么在转成pb的时候,需要加入add_shapes=True
选项,不然的话转出来的pb里面,所有的node信息都不包含_output_shapes
属性。没有这个属性对于卷积层和BN层没有影响,但会影响tvm的转换,按照现在的tvm版本&#