tensorflow訓練的模型在java中的使用

tensorflow訓練模型一般使用python api編寫,簡單記錄下這些模型保存後怎麼在java中調用。java

python中訓練完成,模型保存使用以下api保存:node

# 保存二進制模型  
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模型中,一般是分類輸出的名稱。python

 

例如模型定義時代碼爲:android

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名。git

java中調用代碼片斷: github

public static void main(String[] args) {  
        String labels = "17,16,7,8,3,15,4,14,2,5,12,18,9,10,1,11,13,6";  
  
        TensorFlowInferenceInterface tfi = new TensorFlowInferenceInterface("D:/tf_mode/output_graph.pb","imageType");  
        final Operation operation = tfi.graphOperation("y_conv_add");  
        Output output = operation.output(0);  
        Shape shape = output.shape();  
        final int numClasses = (int) shape.size(1);  
        float[] floatValues = getImagePixel("D:/tf_mode/ci/ci/333.jpg"); //將圖片處理爲輸入對應張量格式  
  
        // 輸入圖片  
        tfi.feed("x_input", floatValues, 1, 2048); //將數據複製給輸入張量x_input即爲模型定義時的x名稱  
        tfi.run(new String[] { "y_conv_add" }, false);//輸出張量  
        float[] outPuts = new float[numClasses];//結果分類  
        tfi.fetch("y_conv_add", outPuts);//接收結果 outPuts保存的即爲預測結果對應的機率,最大的一個一般爲本次預測結果 
}

TensorFlowInferenceInterface參考:api

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

java api和tensorflow的依賴:fetch

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

調用過程參考:

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

相關文章
相關標籤/搜索