java加载tensorflow模型预测

算法同学通常会利用python做tensorflow深度学习模型的训练和测试，但当离线模型训练好之后，往往需要将模型部署上线，将其应用到web端或app应用调用，甚至分布式任务使用等等。
本文将介绍一下，如何利用java加载tensorflow的pb模型，实现模型的在线推理和预测。
JAVA API 参考文档：

https://tensorflow.google.cn/api_docs/java/reference/org/tensorflow/package-summary

注意：TensorFlow Java API 本质上是用Java封装了C++的动态库，而且Java API不在 TensorFlow API 稳定性保障的涵盖范围内。

1、环境搭建

目前java引擎的开发通常会使用maven进行jar包管理，首先需要在pom.xml中添加以下依赖以加载对应jar包。
Ps：至于java 的tensorflow的版本，要和python训练时用的tensorflow的版本一致。
depandency方法1：

<!-- https://mvnrepository.com/artifact/org.tensorflow/tensorflow -->
<dependency>
    <groupId>org.tensorflow</groupId>
    <artifactId>tensorflow</artifactId>
   <version>1.15.0</version>
</dependency>

depandency方法2：

<dependency>
    <groupId>org.tensorflow</groupId>
    <artifactId>libtensorflow</artifactId>
    <version>1.15.0</version>
</dependency>
<dependency>
    <groupId>org.tensorflow</groupId>
    <artifactId>proto</artifactId>
    <version>1.15.0</version>
</dependency>
<dependency>
    <groupId>org.tensorflow</groupId>
    <artifactId>libtensorflow_jni</artifactId>
    <!--artifactId>libtensorflow_jni_gpu</artifactId-->
    <version>1.15.0</version>
</dependency>

2、Java加载tensorflow模型

/**
 * @author lzhenboy
 * @date 2020-01-02
 */
import org.ansj.domain.Term;
import org.tensorflow.SavedModelBundle;
import org.tensorflow.Session;
import org.tensorflow.Tensor;

import java.nio.charset.StandardCharsets;
import java.util.List;

/**
 * Tensorflow pb模型加载 & 预测
 */
public class TfModelLoader {
    private final Session session;

    public TfModelLoader(String modelPath, String modelTag) {
        SavedModelBundle modelBundle = SavedModelBundle.load(modelPath, modelTag);
        this.session = modelBundle.session();
    }

    public TfModelLoader(String modelPath) {
        SavedModelBundle modelBundle = SavedModelBundle.load(modelPath, "serve");
        this.session = modelBundle.session();
    }

    public float[] predict(List<Term> terms) {
        int termSize = terms.size();

        byte[][][] wordsBytes = new byte[1][termSize][];
        int[] nwords = new int[]{wordsBytes[0].length};
        for (int i = 0; i < termSize; ++i) {
            byte[] nameBytes = terms.get(i).getName().getBytes(StandardCharsets.UTF_8);
            wordsBytes[0][i] = nameBytes;
        }

        try (Tensor<?> wordsTensor = Tensor.create(wordsBytes);
             Tensor<?> nwordsTensor = Tensor.create(nwords)) {
            Tensor<?> result = session.runner()
                    .feed("words", wordsTensor)
                    .feed("nwords", nwordsTensor)
                    .fetch("predict").run().get(0);
            float[][] preds = new float[1][wordsBytes[0].length];
            result.copyTo(preds);

            // Tensor Close
            result.close();
            return preds[0];
        } catch (Throwable e) {
            log.error("Model Predict Exception: ", e);
        }
        return null;
    }
}

需要注意的是，参与计算的tensor都要close，以防止内存泄漏（tensorflow底层代码是c++实现，内存管理不受java控制，因此需要人工close）。

3、使用saved_model_cli

如果你使用的是别人的模型，或者你想通过pb模型文件快速查看模型的输出输出标签，saved_model_cli提供了通过命令行检查并恢复模型的快捷方式，可以通过如下方式检查该模型的相关信息：

saved_model_cli show --dir=graph_dir --all
eg:
saved_model_cli show --dir ./saved_model/1604666267 --all

结果如下：

saved_model_cli.png
从上图可以看出：
模型的tag为serve
模型input_placeholder为：(1) word (2) nword
模型output为：logits

参考文献

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