python - Tensorflow keras model.save 使用自定义层引发 InaccessibleTensorError

Question

我有一个自定义 keras 模型（使用tf.keras.model/functional api），我需要为其保存权重和模型架构（使用keras.Model.save）。

该模型使用标准 keras 组件，如密集、扁平等和一个特定的自定义层：

class NormalizationLayer(tf.keras.layers.Layer):

    def __init__(self, name=None, axis=1):
        super(NormalizationLayer, self).__init__(name=name)
        self.max = None
        self.axis = axis

    def build(self, input_shape):
        pass

    def call(self, inputs, inverse=False):
        if inverse:
            return tf.transpose(tf.math.multiply(tf.transpose(inputs), self.max))
        else:
            self.max = tf.reduce_max(inputs, axis=self.axis)
            if self.axis == 1:
                return tf.transpose(tf.math.divide_no_nan(tf.transpose(inputs), self.max))
            else:
                return tf.math.divide_no_nan(inputs, self.max)

    def get_config(self):
        custom_config = {'max': self.max,
                         'axis': self.axis}
        base_config = super(NormalizationLayer, self).get_config()
        base_config.update(custom_config)
        return base_config

get_config当我首先遇到问题时，我已经添加了该方法。不幸的是，在尝试使用 保存模型时model.save，我遇到了以下错误：

tensorflow.python.framework.errors_impl.InaccessibleTensorError: The tensor 'Tensor("Max:0", shape=(None,), dtype=float32)' cannot be accessed here: it is defined in another function or code block. Use return values, explicit Python locals or TensorFlow collections to access it. Defined in: FuncGraph(name=input_normalization_layer_call_and_return_conditional_losses, id=140460908747216); accessed from: FuncGraph(name=model_layer_call_and_return_conditional_losses, id=140460638991312)

我的方法总体上是错误的，还是我监督了什么。已经尝试添加@tf.function-decorator，但遇到了类似的问题。model.save_weights()按预期保存权重。

我正在使用 tensorflow==2.2.0rc4。

感谢您提前提供任何帮助评论。

标签： pythontensorflowkerastensorflow2.x