使用@tf.custom_gradient的自定义Layer报错：仅即时执行支持关键字参数

这个问题我之前也踩过坑，核心原因是哪怕你设置了run_eagerly=True，模型构建的某些环节依然会触发图模式，而在图模式下@tf.custom_gradient装饰器完全不支持关键字参数传递——这也是你单独调用函数时正常，但放到Keras Layer里就报错的根本原因。

下面给你几个优先级从高到低的解决方案：

1. 直接改用位置参数调用自定义梯度函数 #

检查你在Custom_Layer的call方法里调用custom_function的代码，把所有关键字参数改成位置参数传递。

比如，如果你之前是这么写的：

# 错误示例：使用关键字参数触发报错
output = custom_function(inputs, alpha=self.alpha)

改成：

# 正确示例：用位置参数传递
output = custom_function(inputs, self.alpha)

同时确保自定义梯度函数的参数定义也匹配位置参数的逻辑（不要给参数设默认值后又用关键字传递）：

@tf.custom_gradient
def custom_function(x, alpha):  # 这里不要写 alpha=0.5 这种默认值
    # 前向计算逻辑
    y = x * alpha
    # 反向梯度定义
    def grad(dy):
        # 返回对应输入参数的梯度，alpha作为超参数梯度设为None
        return dy * alpha, None
    return y, grad

2. 用functools.partial提前绑定参数 #

如果你的自定义函数必须保留关键字参数的灵活性，可以用functools.partial在Layer初始化阶段就把参数绑定好，避免调用时传递关键字：

from functools import partial

class Custom_Layer(Layer):
    def __init__(self, alpha=0.5, **kwargs):
        super().__init__(**kwargs)
        self.alpha = tf.Variable(alpha, dtype=tf.float32)
        # 提前绑定alpha参数，生成一个无需关键字的新函数
        self.bound_custom_func = partial(custom_function, alpha=self.alpha)
    
    def call(self, inputs):
        # 直接调用绑定好的函数，不用额外传参
        return self.bound_custom_func(inputs)

3. 升级TF版本（可选） #

TF2.0.0属于早期版本，确实存在一些自定义梯度相关的小bug，如果上面两个方法都没解决问题，可以尝试升级到TF2.0.x的稳定版（比如2.0.4），不过这个不是必须的，前两个方法基本能覆盖绝大多数场景。

完整可运行示例 #

给你一个可以直接跑通的测试代码，你可以参考调整自己的实现：

import tensorflow as tf
from tensorflow.keras.layers import Layer, Input
from tensorflow.keras.models import Model

# 定义带自定义梯度的函数
@tf.custom_gradient
def custom_function(x, alpha):
    output = x * alpha
    def grad(dy):
        return dy * alpha, None
    return output, grad

class Custom_Layer(Layer):
    def __init__(self, alpha=0.5, **kwargs):
        super(Custom_Layer, self).__init__(**kwargs)
        self.alpha = tf.Variable(alpha, dtype=tf.float32)
    
    def call(self, inputs):
        # 用位置参数调用自定义函数
        return custom_function(inputs, self.alpha)

# 构建并编译模型
input_length = 10
input_tensor = Input(shape=(1, input_length))
output_layer = Custom_Layer(alpha=0.3)(input_tensor)
model = Model([input_tensor], [output_layer])
model.compile(optimizer='rmsprop', run_eagerly=True, loss='mae', metrics=['accuracy'])
model.summary()

# 测试训练流程
import numpy as np
x_test = np.random.rand(32, 1, input_length)
y_test = np.random.rand(32, 1, input_length)
loss, acc = model.train_on_batch(x_test, y_test)
print(f"训练损失: {loss}, 准确率: {acc}")

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