能否用Numba子类化Keras Layer？训练报错如何解决？

这个问题我之前也碰到过——eager模式下跑起来没问题，一到训练阶段用model.fit就报错，核心原因就是tf.function构建计算图时不允许直接调用.numpy()，因为这会把张量从图里“拉”出来，破坏了TensorFlow的图追踪机制。

咱们来一步步修改代码，让Numba和Keras的图模式完美兼容：

核心思路 #

不要在tf.custom_gradient里手动提取numpy数组，而是用TensorFlow官方的tf.numpy_function作为中间层，它能安全地在计算图和Numba编译的numpy函数之间转换数据，而且完全兼容图模式。

修改后的完整代码 #

import numpy as np
import numba
import tensorflow as tf

# 保留你的Numba编译函数，这部分不需要改
@numba.jit(nopython=True)
def func(param, input):
    return param * input**2

@numba.jit(nopython=True)
def gradfunc(param, input):
    return input**2

# 重写自定义梯度函数，完全兼容图模式
@tf.custom_gradient
def func_tf(param, input):
    # 用tf.numpy_function调用Numba函数，自动处理张量和numpy数组的转换
    result = tf.numpy_function(func, [param, input], tf.float32)
    # 手动设置结果形状，因为tf.numpy_function会丢失形状信息
    result.set_shape(input.shape)
    
    def grad(dy):
        # 对param的梯度：上游梯度dy * gradfunc的结果
        # 用lambda包装，把dy、param、input传给Numba的gradfunc
        grad_param = tf.numpy_function(
            lambda p, i, dy_val: gradfunc(p, i) * dy_val,
            [param, input, dy],
            tf.float32
        )
        grad_param.set_shape(param.shape)
        
        # 对input的梯度：用TensorFlow原生操作实现，高效且兼容图模式
        grad_input = dy * 2 * param * input
        
        return grad_param, grad_input
    
    return result, grad

class myLayer(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()
    def build(self, input_shape):
        # 给参数设置初始值和形状，避免后续报错
        self.param = self.add_weight("param", initializer="ones", shape=())
    def call(self, input):
        return func_tf(self.param, input)

class myModel(tf.keras.Model):
    def __init__(self, num_layers):
        super().__init__(name='')
        self._layers = [myLayer() for _ in range(num_layers)]
    def call(self, input_tensor):
        for layer in self._layers:
            input_tensor = layer(input_tensor)
        return input_tensor

# 测试eager模式
model = myModel(3)
print(model(1.5))

# 训练测试
def loss(target, output):
    return tf.abs(tf.reduce_sum(target - output))**2

model.compile(
    optimizer=tf.keras.optimizers.Adam(),
    loss=loss,
    metrics=[loss])

# 现在可以正常训练了
history = model.fit([0.1], [0.4], batch_size=None, epochs=5)

关键修改点说明 #

1. 移除.numpy()调用：
   原来的func_tf里直接用param.numpy()和input.numpy()提取数值，这在图模式下是禁止的。改用tf.numpy_function后，TensorFlow会在图运行时自动完成张量到numpy数组的转换，不需要手动干预。

2. 恢复形状信息：
   tf.numpy_function返回的张量会丢失形状信息，所以必须用set_shape()手动设置，否则Keras在构建模型时会因为形状不匹配报错。

3. 梯度计算的链式法则：
   在梯度函数里，上游梯度dy需要和Numba计算的局部梯度相乘，这里用lambda函数把dy、param、input一起传给Numba的gradfunc，确保梯度计算符合链式法则。

4. 混合使用TensorFlow原生操作：
   对于像2*param*input这种简单的梯度计算，直接用TensorFlow原生操作更高效，也避免了额外的Numba调用开销。如果你的梯度逻辑复杂，也可以用同样的tf.numpy_function方式调用Numba函数。

这样修改后，不管是eager模式还是训练时的图模式，代码都能正常运行，同时保留了Numba带来的计算加速。

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