Keras自定义层结合常量变量实现历史值存储复用问题求助

• 变量更新方式完全错误：call方法里写self.prev_input = inputs_、self.prev_output = output是直接把类属性的指向从原本的tf.Variable持久化对象，改成了当前步的临时计算张量，根本没有对原变量做值写入，跑完一次前向传播历史值根本存不下来。更新Variable的存储值必须用自带的assign方法，不能直接给属性重赋值。
• 常量取值方式不兼容图模式：K.get_value(c)是eager模式下把张量转成numpy值的操作，Keras默认走静态图编译的时候，这个操作会直接报错。本身定义的c就是tf常量，直接放到计算图里参与运算就行，不需要额外取数值。
• get_config实现不符合规范：这个方法只需要返回层初始化时传入的可序列化参数，用来做模型保存/加载的重建，直接把tf.Variable对象塞进config字典，序列化的时候会直接报错。而且prev_input、prev_output是层的内部状态，不是初始化参数，不需要放到config里。
• 输出形状定义不匹配：初始化变量用的是input_dim维度，但是compute_output_shape硬编码返回[None, 1]，只要input_dim不等于1，就会触发维度不匹配的报错。
• 变量形状未适配动态批次：初始化变量时只考虑了特征维度，没有预留批次维度的适配空间，遇到不同batch size的输入时，赋值会触发形状校验错误。

import tensorflow as tf
from tensorflow.keras.layers import Layer

# 定义外部常量
c = tf.constant(5., dtype=tf.float32)

class MyLayer(Layer):
    def __init__(self, input_dim, **kwargs):
        super(MyLayer, self).__init__(**kwargs)
        self.input_dim = input_dim
        # 初始化状态变量，关闭固定形状校验以适配动态批次大小
        self.prev_input = tf.Variable(
            initial_value=tf.zeros((1, input_dim)),
            trainable=False,
            validate_shape=False
        )
        self.prev_output = tf.Variable(
            initial_value=tf.zeros((1, input_dim)),
            trainable=False,
            validate_shape=False
        )

    def get_config(self):
        # 仅返回初始化需要的可序列化参数
        config = super().get_config()
        config.update({
            "input_dim": self.input_dim
        })
        return config

    def call(self, inputs_, training=False):
        a = inputs_[0]
        b = inputs_[1]
        # 常量直接参与计算图运算
        output = self.prev_output * a * b / c
        
        # 通过assign方法更新变量存储的历史值，不修改属性指向
        self.prev_input.assign(inputs_)
        self.prev_output.assign(output)
        
        return output

    def compute_output_shape(self, input_shape):
        # 输出形状与输入特征维度对齐
        return tf.TensorShape([None, self.input_dim])

    def reset_state(self):
        # 新增状态重置方法，需要清空历史值时调用
        self.prev_input.assign(tf.zeros((1, self.input_dim)))
        self.prev_output.assign(tf.zeros((1, self.input_dim)))

• 层的历史值默认是跨批次留存的，如果训练时需要每个epoch/每次推理前清空历史值，主动调用reset_state方法即可，也可以配合Keras回调在对应生命周期节点自动调用。
• 该实现同时兼容eager执行模式和静态图编译模式，不需要修改模型的run_eagerly参数就能正常更新状态。
• 如果输入的特征维度会动态变化，不要在初始化时固定input_dim，可以在build方法里根据第一次传入的输入形状创建状态变量，灵活性更高。

内容的提问来源于stack exchange，提问作者alexka