Python 3.x Tensorflow 2.0:pb用于恢复模型和执行推理

我试图修改官方文件，以便保存生成器的模型，然后能够重新加载并执行图像生成（推断）

为了做到这一点，我遵循了本笔记本中的示例，代码执行得很好，没有任何问题

因此，基本上我在训练循环的末尾添加了以下代码，以将生成器的模型保存到磁盘：

tf.saved_model.save(generator, checkpoint_dir + "/pix2pix/1/")
  print("\nGenerator model saved in {}/pix2pix/1/".format(checkpoint_dir))

然后，为了重新加载模型并进行推断，我使用以下代码：

def generate_images2(model, test_input, tar, debug=True):

  # Key name of the output layer in the model:
  key_name = list(model.structured_outputs.keys())[0]

  prediction = model(test_input)[key_name]
  if debug:
    print("Key name of the output layer in the model:", key_name)
    print("\nprediction:\n", prediction)

  plt.figure(figsize=(15,15))

  display_list = [test_input[0], tar[0], prediction[0]]
  title = ['Input Image', 'Ground Truth', 'Predicted Image']

  for i in range(3):
    plt.subplot(1, 3, i+1)
    plt.title(title[i])
    # getting the pixel values between [0, 1] to plot it.
    plt.imshow(display_list[i] * 0.5 + 0.5)
    plt.axis('off')
  plt.show()

checkpoint_dir = './training_checkpoints'

loaded = tf.saved_model.load(checkpoint_dir + "/pix2pix/1/")
print("Load Model - signature keys:", list(loaded.signatures.keys()))

infer = loaded.signatures["serving_default"]
print("\nInference structure output:", infer.structured_outputs)


# Run the trained model on one exampel of the test dataset
for inp, tar in test_dataset.take(1):
  generate_images2(infer, inp, tar, debug=True)

代码运行时没有错误，但推断步骤的输出张量（基于重新加载的模型）仅包含NaN值。以下是推理的输出：

Load Model - signature keys: ['serving_default']

Inference structure output: {'conv2d_transpose_8': TensorSpec(shape=(None, None, None, 3), dtype=tf.float32, name='conv2d_transpose_8')}

Key name of the output layer in the model: conv2d_transpose_8

prediction:
 tf.Tensor(
[[[[nan nan nan]
   [nan nan nan]
   [nan nan nan]
   ...
   [nan nan nan]
   [nan nan nan]
   [nan nan nan]]

  [[nan nan nan]
   [nan nan nan]
   [nan nan nan]
   ...
   [nan nan nan]
   [nan nan nan]
   [nan nan nan]]

  [[nan nan nan]
   [nan nan nan]
   [nan nan nan]
   ...
   [nan nan nan]
   [nan nan nan]
   [nan nan nan]]

  ...

  [[nan nan nan]
   [nan nan nan]
   [nan nan nan]
   ...
   [nan nan nan]
   [nan nan nan]
   [nan nan nan]]

  [[nan nan nan]
   [nan nan nan]
   [nan nan nan]
   ...
   [nan nan nan]
   [nan nan nan]
   [nan nan nan]]

  [[nan nan nan]
   [nan nan nan]
   [nan nan nan]
   ...
   [nan nan nan]
   [nan nan nan]
   [nan nan nan]]]], shape=(1, 256, 256, 3), dtype=float32)

我是否遗漏了什么，或者Tensorflow 2.0存在问题。有什么想法或建议吗

目前，我怀疑通过将多个Keras顺序模型（作为构建块）组合在一起构建模型的方式可能会导致模型保存/恢复功能出现问题。（类似地，Tensorboard也无法正确显示模型的图形：每个构建块似乎与其他块断开连接）

注:

• 我将Tensorflow GPU 2.0.0-rc1与Python 3.6.8一起使用（与rc0的问题相同）
• 这是一个-这个笔记本可以在不使用Kubeflow集群的情况下重现问题
• 这是一个
• 此笔记本中的型号需要至少16GB内存的GPU