tensorflow训练的模型在java中的使用

本文介绍如何在Java环境中加载并使用TensorFlow的PB模型文件。重点包括Python训练模型后保存为PB格式的方法,以及Java环境下通过TensorFlowInferenceInterface进行模型调用的具体步骤。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

tensorflow训练模型通常使用python api编写,简单记录下这些模型保存后怎么在java中调用。

python中训练完成,模型保存使用如下api保存:

[python]  view plain  copy
  1. # 保存二进制模型  
  2. output_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, output_node_names=['y_conv_add'])  
  3. with tf.gfile.FastGFile('/logs/mnist.pb', mode='wb') as f:  
  4.     f.write(output_graph_def.SerializeToString())  
保存为二进制pb文件,主要的点是output_node_names数组,该数据的名称表示需要保存的tensorflow tensor名。既是在python中定义模型时指定的计算操作的name。填写什么就保存到什么节点。在cnn模型中,通常是分类输出的名称。

例如模型定义时代码为:

[python]  view plain  copy
  1. y_conv = tf.add(tf.matmul(h_fc1_drop, W_fc2), b_fc2, name='y_conv_add'# cnn输出层,名称y_conv_add  
  2. # 训练和评价模型  
  3. softmax = tf.nn.softmax_cross_entropy_with_logits(labels=y_, logits=y_conv)  

模型在java中使用需要关心模型输入tensor和输出tensor名,所以定义模型时,所有的输入tensor最好指定名称,如输入x和dropout名。

java中调用代码片段:

[java]  view plain  copy
  1. public static void main(String[] args) {  
  2.         String labels = "17,16,7,8,3,15,4,14,2,5,12,18,9,10,1,11,13,6";  
  3.   
  4.         TensorFlowInferenceInterface tfi = new TensorFlowInferenceInterface("D:/tf_mode/output_graph.pb","imageType");  
  5.         final Operation operation = tfi.graphOperation("y_conv_add");  
  6.         Output output = operation.output(0);  
  7.         Shape shape = output.shape();  
  8.         final int numClasses = (int) shape.size(1);  
  9.         float[] floatValues = getImagePixel("D:/tf_mode/ci/ci/333.jpg"); //将图片处理为输入对应张量格式  
  10.   
  11.         // 输入图片  
  12.         tfi.feed("x_input", floatValues, 12048); //将数据复制给输入张量x_input即为模型定义时的x名称  
  13.         tfi.run(new String[] { "y_conv_add" }, false);//输出张量  
  14.         float[] outPuts = new float[numClasses];//结果分类  
  15.         tfi.fetch("y_conv_add", outPuts);//接收结果 outPuts保存的即为预测结果对应的概率,最大的一个通常为本次预测结果  
[java]  view plain  copy
  1. }  

TensorFlowInferenceInterface参考:

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/android/java/org/tensorflow/contrib/android/TensorFlowInferenceInterface.java

java api和tensorflow的依赖:

https://github.com/tensorflow/tensorflow/tree/master/tensorflow/java

调用过程参考:

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/android/src/org/tensorflow/demo/TensorFlowImageClassifier.java

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值