Java加载TensorFlow PB模型打包Jar后报Invalid GraphDef错误

这个报错的核心原因是你读取模型文件的方式有问题！当模型文件在Jar包内部时，InputStream.available()并不能准确返回文件的完整大小，而且is.read(graphDef)也无法保证一次性读取完所有字节，这就导致你得到的graphDef是不完整的，TensorFlow自然无法解析出有效的GraphDef。

正确的资源读取方式 #

你需要用ByteArrayOutputStream来完整读取输入流的所有字节，而不是依赖available()和单次read()调用。修改后的代码如下：

InputStream is = LstmSeqModel.class.getResourceAsStream("/lstmSeq.pb");
ByteArrayOutputStream baos = new ByteArrayOutputStream();
byte[] buffer = new byte[1024];
int length;
while ((length = is.read(buffer)) != -1) {
    baos.write(buffer, 0, length);
}
byte[] graphDef = baos.toByteArray();
g.importGraphDef(graphDef);
sess = new Session(g);
// 别忘了关闭流
is.close();
baos.close();

额外注意事项 #

• 确保打包Jar时，lstmSeq.pb确实被正确打包进去了。在IntelliJ IDEA中，检查你的resources目录是否被标记为Resources Root，这样打包时才会将目录下的文件包含到Jar的根路径下。
• 如果模型文件很大，建议使用缓冲流来提升读取效率，比如用BufferedInputStream包裹原输入流：

  InputStream is = new BufferedInputStream(LstmSeqModel.class.getResourceAsStream("/lstmSeq.pb"));
  

这样修改后，无论模型在本地还是Jar包中，都能完整读取模型文件的所有字节，就不会再出现Invalid GraphDef的报错了。

内容的提问来源于stack exchange，提问作者尉思谜