修改tf.keras.Optimizer参数触发TypeError的解决方法求助

我尝试实现一个参考他人思路的Adam优化器梯度累积包装器AccumOptimizer，代码如下：

class AccumOptimizer(tf.keras.optimizers.Optimizer):

    def __init__(self, optimizer, steps_per_update=1, **kwargs):

        super(AccumOptimizer, self).__init__(name="AccumOptimizer", **kwargs)
        self.optimizer = optimizer
        self.steps_per_update = steps_per_update
        self.iterations = tf.Variable(0, dtype='int64', name='iterations')
        self.cond = tf.equal(self.iterations % self.steps_per_update, 0)
        self.lr = self.optimizer.learning_rate
        self.optimizer.learning_rate = tf.cond(self.cond,
                                               lambda: self.optimizer.learning_rate,
                                               lambda: tf.constant(0, tf.float32))
        ...

训练时触发如下TypeError：

TypeError: __array__() takes 1 positional argument but 2 were given

报错指向通过self.cond更新self.optimizer.learning_rate的代码段，我确认tf.cond在分支返回单值时应返回单个结果。受限于使用TensorFlow 1.15.x版本，请问该如何解决这个问题？

在TensorFlow 1.15.x中，直接将tf.cond的结果赋值给优化器的learning_rate会引发类型兼容问题——优化器的学习率属性期望可调用对象或静态张量，而tf.cond返回的动态计算张量在后续内部数组转换时会触发参数不匹配错误。以下是两种可行的解决方式：

方法一：直接控制梯度更新时机（推荐） #

移除__init__中修改self.optimizer.learning_rate的代码，转而在梯度累积逻辑中，仅当满足累积步数条件时才调用原优化器执行更新，不满足时跳过：

class AccumOptimizer(tf.keras.optimizers.Optimizer):
    def __init__(self, optimizer, steps_per_update=1, **kwargs):
        super(AccumOptimizer, self).__init__(name="AccumOptimizer", **kwargs)
        self.optimizer = optimizer
        self.steps_per_update = steps_per_update
        self.iterations = tf.Variable(0, dtype='int64', name='iterations')
        # 初始化梯度累积变量
        self.accum_grads = []
        for var in tf.trainable_variables():
            self.accum_grads.append(tf.Variable(tf.zeros_like(var), trainable=False))

    def apply_gradients(self, grads_and_vars, name=None):
        # 累积当前步梯度
        for (grad, var), accum_grad in zip(grads_and_vars, self.accum_grads):
            accum_grad.assign_add(grad)
        
        # 更新迭代计数
        self.iterations.assign_add(1)
        
        # 定义满足条件时的更新操作：应用累积梯度+重置累积器
        def apply_update():
            update_op = self.optimizer.apply_gradients(
                [(accum_grad, var) for (_, var), accum_grad in zip(grads_and_vars, self.accum_grads)]
            )
            reset_ops = [accum_grad.assign(tf.zeros_like(accum_grad)) for accum_grad in self.accum_grads]
            return tf.group(update_op, *reset_ops)
        
        # 不满足条件时执行空操作
        def skip_update():
            return tf.no_op()
        
        # 根据条件分支执行对应逻辑
        return tf.cond(
            tf.equal(self.iterations % self.steps_per_update, 0),
            apply_update,
            skip_update
        )

方法二：包装学习率为可调用函数 #

如果需要通过控制学习率实现逻辑，可以将学习率包装为可调用的lambda函数，避免直接赋值tf.cond的结果：

class AccumOptimizer(tf.keras.optimizers.Optimizer):
    def __init__(self, optimizer, steps_per_update=1, **kwargs):
        super(AccumOptimizer, self).__init__(name="AccumOptimizer", **kwargs)
        self.optimizer = optimizer
        self.steps_per_update = steps_per_update
        self.iterations = tf.Variable(0, dtype='int64', name='iterations')
        self.base_lr = self.optimizer.learning_rate
        # 将学习率包装为可调用函数，动态返回对应值
        self.optimizer.learning_rate = lambda: tf.cond(
            tf.equal(self.iterations % self.steps_per_update, 0),
            lambda: self.base_lr,
            lambda: tf.constant(0.0, dtype=tf.float32)
        )
        # 其他初始化逻辑...

注意：这种方式需要确保原优化器在TF1.15中支持接收可调用的学习率对象，部分早期Keras优化器可能存在适配问题，因此更推荐方法一的梯度累积逻辑。

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