无法将自定义模型保存为.pb格式（tensorflow 2.1.0）

Question

我正在尝试从具有自定义层的 Keras 模型 (.h5/json) 转换为张量流模型 (.pb)，并且我使用以下自定义 LSTM 层。

我可以使用自定义层加载模型，也可以在 tensorflow.keras 上使用它。但是，当我尝试时遇到以下问题loaded_model.save(model_name, save_format="tf")

Traceback (most recent call last):

  File "/home/usr7/s70087f/.local/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 1619, in _create_c_op

    c_op = c_api.TF_FinishOperation(op_desc)

tensorflow.python.framework.errors_impl.InvalidArgumentError: Shape must be rank 2 but is rank 1 for 'MatMul_7' (op: 'MatMul') with input shapes: [256], [256,256].

V_tm1_i（带有Attention的LSTM输入门的计算）有一个问题：

    i = self.recurrent_activation(

        x_i 

        + K.dot(V_tm1_i, self.recurrent_kernel[:, :self.units])

        + K.dot(c_i, self.context_kernel[:, :self.units]))

我知道我们必须使用 >= 2 ndims 的张量来“MatMul”，但实际上它有一个批量大小维度。因此我希望它成为[None,256], [256,256]。

显示自定义层的输入和隐藏状态：

AttentionVLSTM

inputs

Tensor("model/Decoder_Attention_VLSTM/strided_slice_1:0", shape=(None, 256), dtype=float32)

states

(<tf.Tensor 'model/reshape/Reshape:0' shape=(None, 27136) dtype=float32>, <tf.Tensor 'model/Decoder_Activation_2/Relu:0' shape=(None, 256) dtype=float32>)

AttentionVLSTM

inputs

Tensor("TensorArrayV2Read/TensorListGetItem:0", shape=(None, 256), dtype=float32)

states

(<tf.Tensor 'Placeholder_3:0' shape=(None, 27136) dtype=float32>, <tf.Tensor 'Placeholder_4:0' shape=(None, 256) dtype=float32>)

AttentionVLSTM

inputs

Tensor("inputs:0", shape=(None, 256), dtype=float32)

states

(<tf.Tensor 'states:0' shape=(27136,) dtype=float32>, <tf.Tensor 'states_1:0' shape=(256,) dtype=float32>)

V_tm1是从隐藏状态获得的，“Placeholder_4:0”是正确的形状[None,256]，但“states_1:0”丢失了batch axis的维度[256,]。其他自定义层（除了本层）可以运行，并且该层也在tensorflow.Keras中运行。

Answer 1

我可以解决这个问题。错误的原因是该自定义类末尾的初始状态的定义。这是不必要的，因为初始状态已经在“AbstractRNNCell”中定义。