Tensorflow model.save() 不保存模型的推理函数

Question

我有一个类，它扩展了keras.Model，并实现了一个由几个标准层和自定义层组成的模型。我不是使用model.fit()，而是使用for循环来迭代数据并运行以下训练步骤函数

@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
inputs = batch[0] + [batch[1][-1]]
predictions, _ = model(inputs, training=True)
loss = log_likelihood(batch[1], predictions, batch[2])
regularisation_loss = model.output_layers.losses
gradients = tape.gradient([regularisation_loss, loss], model.trainable_variables)
optimiser.apply_gradients(zip(gradients, model.trainable_variables))

为了节省开支，我只是简单地调用model.save(model_path)。在我尝试重新加载保存的模型之前，一切都很顺利。加载调用抛出(仅底线(

File "...condaenvstf.2libsite-packagestensorflowpythonsaved_modelfunction_deserialization.py", line 265, in recreate_function
concrete_function_objects.append(concrete_functions[concrete_function_name])
KeyError: '__inference_model_layer_call_fn_37936'

我认为，这是指我的模型的主要call函数，它只是由于某种原因而没有序列化和保存。我查看了concrete_functions，除了这个函数外，几乎所有的东西都在里面。我在使用和不使用tf.function装饰器的情况下都进行了尝试，但没有发现任何区别。我现在有点迷路了，我真的很感激不仅有一个答案，还有一个很好的调试方向。

感谢

Answer 1

因此，在https://github.com/tensorflow/tensorflow/issues/42004这让我走上了正确的方向，我做了以下的改变，似乎已经解决了的问题

1. 除了"；主模型包装器thingy"；，是对图层对象的更改
2. 所有自定义实现(层和模型(都被赋予了get_config函数

我不确定是哪一个成功了，但我的钱在第一点上。