Java加载tensorflow模型(单机加载或分布式加载)

我们常用python进行tensorflow深度模型训练，然后，训练后的模型需要应用到web端调用，或者app应用调用，甚至分布式任务使用等等。在这些应用中java代码调用，是避免不了的。本文就介绍一下，java加载tensorflow模型的方式，分别单机调用和分布式调用。

1.环境需要

首先需要加载一些jar包，如果是maven项目在pom.xml中添加以下依赖。至于版本信息，切记一定要和你python代码训练的时候的tensorflow版本一致。

<dependency>
    <groupId>org.tensorflow</groupId>
    <artifactId>tensorflow</artifactId>
    <version>0.13</version>
</dependency>

2.java单机加载tensorflow模型

单价加载模型的代码简要如下，做了简化哈，但是核心的都在这里。offLineInit函数就是单机或者叫离线的模型加载方式。这里的模型一般是pb模型。可以是序列化的也可以是savemodel形式的。

import org.tensorflow.Session;
import org.tensorflow.Graph;
import org.tensorflow.Tensor;
import java.io.File;
import java.io.BufferedInputStream;
import java.io.FileInputStream;

public class SplitByDeepModel  implements Serializable{
	
	public Session session;
	
    public boolean offLineInit(String path) throws Exception{
    	/*
    	 * 单机加载tensorflow模型，初始化环境
    	 */
    	if (!this.tensorflow_factory.InitModelsGeneratorByBasePath(path))
			throw new UDFArgumentException("tensorflow model Generator Init failed");			
        byte[] graphDef = readAllBytes(path+"/models/model_test1.pb");
        Graph graph = new Graph();
        graph.importGraphDef(graphDef);
        session = new Session(graph);
    	return true;
    }
}

3.分布式tensorflow模型加载

分布式加载方式如下，代码如下：

public boolean InitModelByBasePath(String tensorflow_model_path){
    	try {
        	// model label
        	modelTag  = "Mytest";
        	modelPath = tensorflow_model_path;
			System.err.println("tensorflow models Path:" + modelPath);
			
        	Configuration conf = new Configuration();
        	Graph graph = new Graph();
        	FileSystem fs = FileSystem.get(conf);
        	ByteArrayOutputStream contents = new ByteArrayOutputStream();
        	IOUtils.copy(fs.open(new Path(modelPath)), contents);
        	graph.importGraphDef(contents.toByteArray());
        	this.session = new Session(graph);
    		
    	}catch(Exception e) {
			System.err.println("Init tensorflow model from " + modelPath + " error: " + e.toString());
			e.printStackTrace();
		}
    	return false;
    }    

其中，tensorflow_model_path，hdfs模型路径。这种方式需要将模型文件加载到hdfs系统中，然后应用在每个executor里面进行加载。

3.分布式加载和单机加载的差异

差异主要在文件读取的方式上，其他都是一样的，都是sess加载模型图，然后初始化。
这里有一点需要强调一下。无论是单机模式还是分布式模式加载，都是使用的pb模型文件，加载之后相当于一个函数。但是这个函数是非序列化的，不能跨节点传参。这要注意。。

4.模型加载后使用
代码如下：

    public float[][] lstmPredict(int[][] input) {
    	 Tensor input_X = Tensor.create(input);
    	 int[] seqs = new int[1];
    	 seqs[0] = input[0].length;
    	 Tensor seq_len = Tensor.create(seqs);
    	 Tensor dropout = Tensor.create(1.0f);    	 
         Tensor out= this.session.runner().feed("word_ids",input_X).feed("sequence_lengths", seq_len).
        		 feed("dropout", dropout).fetch("proj/predict").run().get(0);
         
         float [][][] ans = new float[1][input[0].length][3];
         out.copyTo(ans);
         
         // tensor对象必须close对象，来清理内存，因为，tensorflow的底层代码都是c++写的
         // c++的内存释放不受java管理，所以必须主动释放，否做很容易出现内存溢出问题
         input_X.close();
         out.close();
         seq_len.close();
         dropout.close();
         return ans[0];
    }

其中Tensor变量都是，预测时候必须输入的变量以及输出的变量。
这个代码是一个lstm的case。其中：seq_len，dropout，input_X都是必须的输出；out是输出形式。

这里特别强调下，最后的close部分，一定要每次预测之后，所有参与计算的tensor都要close。
为什么？因为tensorflow的底层代码都是c++写的，它的内存管理不受java控制，所以必须手工close，否则很容易出现内存溢出的问题，如果是分布式跑，很可能服务器都宕机了。

我在做这件事的时候，也是经历了各种趟坑，然后在一位大神的指导下，实现了分布式跑预测模型。这也是小白必须经历的过程。