算法同學(xué)通常會(huì)利用python做tensorflow深度學(xué)習(xí)模型的訓(xùn)練和測(cè)試，但當(dāng)離線模型訓(xùn)練好之后，往往需要將模型部署上線，將其應(yīng)用到web端或app應(yīng)用調(diào)用，甚至分布式任務(wù)使用等等。
本文將介紹一下，如何利用java加載tensorflow的pb模型，實(shí)現(xiàn)模型的在線推理和預(yù)測(cè)。
JAVA API 參考文檔：

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

注意：TensorFlow Java API 本質(zhì)上是用Java封裝了C++的動(dòng)態(tài)庫(kù)，而且Java API不在 TensorFlow API 穩(wěn)定性保障的涵蓋范圍內(nèi)。

1、環(huán)境搭建

目前java引擎的開(kāi)發(fā)通常會(huì)使用maven進(jìn)行jar包管理，首先需要在pom.xml中添加以下依賴(lài)以加載對(duì)應(yīng)jar包。
Ps：至于java 的tensorflow的版本，要和python訓(xùn)練時(shí)用的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模型加載 & 預(yù)測(cè)
 */
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;
    }
}

需要注意的是，參與計(jì)算的tensor都要close，以防止內(nèi)存泄漏（tensorflow底層代碼是c++實(shí)現(xiàn)，內(nèi)存管理不受java控制，因此需要人工close）。

3、使用saved_model_cli

如果你使用的是別人的模型，或者你想通過(guò)pb模型文件快速查看模型的輸出輸出標(biāo)簽，saved_model_cli提供了通過(guò)命令行檢查并恢復(fù)模型的快捷方式，可以通過(guò)如下方式檢查該模型的相關(guān)信息：

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

結(jié)果如下：

saved_model_cli.png

參考文獻(xiàn)

Java 版 TensorFlow
Java加載tensorflow模型(單機(jī)加載或分布式加載)