算法同學(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
中添加以下依賴以加載對(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é)果如下:
從上圖可以看出:
模型的
tag
為serve模型
input_placeholder
為:(1) word
(2) nword
模型
output
為:logits