tensorflow.keras 在 python 中保存模型并在 Java 中加载

Question

我有一个经过微调的 vgg 模型，我使用 tensorflow.keras 功能 API 创建了模型，并使用 tf.contrib.saved_model.save_keras_model 保存了模型。 因此模型以这种结构保存：包含 saved_model.json 文件、saved_model.pb 文件的 assets 文件夹和包含 checkpoint 的 variables 文件夹strong>、variables.data-00000-of-00001 和 variables.index。

我可以轻松地在 python 中加载我的模型并使用 tf.contrib.saved_model.load_keras_model(saved_model_path) 获得预测，但我不知道如何在 JAVA 中加载模型。我google了很多，发现这个How to export Keras .h5 to tensorflow .pb?导出为pb文件，然后按照这个链接Loading in Java加载它。我无法冻结图形，并且我尝试使用 simple_save，但 tensorflow.keras 不支持 simple_save（AttributeError: module 'tensorflow.contrib.saved_model' 没有属性 'simple_save'）。那么有人可以帮我弄清楚在JAVA中加载我的模型（tensorflow.keras功能API模型）需要哪些步骤。

我拥有的 saved_model.pb 文件是否足以在 JAVA 端加载？我需要创建输入/输出占位符吗？那怎么导出呢？
感谢您的帮助。

Answer 1

如果您有一个以 SavedModel 格式保存的模型（您似乎已经这样做了，tf.contrib.saved_model.save_keras_model 之类的东西可以帮助创建），那么在 Java 中您可以使用 SavedModelBundle.load 来加载和提供它。您不需要“冻结”模型。

您可能会发现以下有用：

• Example program using the SavedModel format for an object detection model in Java
• Slide deck 关于 Java API 和 SavedModel 格式

但基本思想是您的代码将类似于：

try (SavedModelBundle model = SavedModelBundle.load("<directory>", "serve")) {
  try (Tensor<?> input = makeInputTensor();
       Tensor<?> output = model.session().runner().feed("INPUT_TENSOR", input).fetch("OUTPUT_TENSOR", output).run().get(0)) {
  // Use output
  }
}

其中"INPUT_TENSOR" 和"OUTPUT_TENSOR" 是TensorFlow 图中输入和输出节点的名称。安装 TensorFlow for Python 时安装的 saved_model_cli 命令行工具可以显示模型中这些张量的名称。

请注意，使用 TensorFlow Java API 可能比其他评论者建议的使用 TensorFlow Lite 更适合服务器/桌面应用程序。这是因为 TensorFLow Lite 运行时虽然针对小型设备进行了优化（在内存占用等方面），但还不能导出所有模型。而 TensorFlow Java API 使用完全相同的运行时，因此具有与 TensorFlow for Python 完全相同的功能。

希望对您有所帮助。