tensorflow训练模型通常使用python api编写,简单记录下这些模型保存后怎么在java中调用。
python中训练完成,模型保存使用如下api保存:
# 保存二进制模型
output_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, output_node_names=['y_conv_add'])
with tf.gfile.FastGFile('/logs/mnist.pb', mode='wb') as f:
f.write(output_graph_def.SerializeToString())
保存为二进制pb文件,主要的点是output_node_names数组,该数据的名称表示需要保存的tensorflow tensor名。既是在python中定义模型时指定的计算操作的name。填写什么就保存到什么节点。在cnn模型中,通常是分类输出的名称。
例如模型定义时代码为:
y_conv = tf.add(tf.matmul(h_fc1_drop, W_fc2), b_fc2, name='y_conv_add') # cnn输出层,名称y_conv_add
# 训练和评价模型
softmax = tf.nn.softmax_cross_entropy_with_logits(labels=y_, logits=y_conv)
模型在java中使用需要关心模型输入tensor和输出tensor名,所以定义模型时,所有的输入tensor最好指定名称,如输入x和dropout名。
ja