java 打包成jar包,加载tensorflow模型报错

今天遇到一个问题,没有解决, mark一下。

java加载tensorflow的pb模型时,传递模型所在路径加载没有问题,想着放在resource中加载试一下,然后在本地运行依旧没有问题,但是打包成jar包,再运行就报错:

Exception in thread "main" java.lang.IllegalArgumentException: Invalid GraphDef
at org.tensorflow.Graph.importGraphDef(Native Method)
at org.tensorflow.Graph.importGraphDef(Graph.java:130)
at org.tensorflow.Graph.importGraphDef(Graph.java:114)


没有找到解决办法,谁知道请留言,谢谢。

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

解决了!

原因是采用InputStream读取模型再转成byte[]时数据有丢失。解决办法就是100个字节100个字节这样子读,存在输出流中,再转成byte[]。

贴代码:

InputStream is = LstmSeqModel.class.getResourceAsStream("/lstmSeq.pb");
ByteArrayOutputStream swapStream = new ByteArrayOutputStream();
byte[] buff = new byte[100]; //每次读100个字节
int len = 0;
while ((len = is.read(buff, 0, 100)) > 0){
      swapStream.write(buff, 0, len);
      }
graphDef = swapStream.toByteArray();
g.importGraphDef(graphDef);

猜你喜欢

转载自blog.csdn.net/mimi9919/article/details/79578707