自定义Keras层中调用tf.py_func执行Python函数报错，求正确实现

这个问题我之前也遇到过！在TF2.x版本里，tf.py_func其实已经被弃用了，而且它本身对张量形状的处理很不友好，这大概率是你触发那个TypeError的原因——因为TF没法推断自定义函数输出的形状，导致后续计算中出现了NoneType的形状值，和整数做运算就报错了。

下面给你讲在自定义Keras层里调用Python函数的正确姿势，分两种常用场景：不需要自动微分，和需要支持训练（自动微分）的情况。

场景1：不需要自动微分，仅做自定义计算 #

步骤1：编写Python函数

这个函数要接收numpy数组、返回numpy数组，不要直接处理TF张量：

import tensorflow as tf
import numpy as np

def my_custom_func(input_np):
    # 这里写你的自定义逻辑，比如对每个元素做平方加1
    processed = input_np ** 2 + 1
    # 确保返回类型和后续TF张量匹配，比如float32
    return processed.astype(np.float32)

步骤2：自定义Keras层

继承tf.keras.layers.Layer，用tf.numpy_function替代tf.py_func，核心是必须显式指定输出形状：

class MyCustomLayer(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    
    def call(self, inputs):
        # tf.numpy_function是TF2.x推荐的替代方案
        output = tf.numpy_function(
            func=my_custom_func,
            inp=[inputs],  # 输入张量的列表
            Tout=tf.float32  # 输出张量的类型
        )
        # 手动设置输出形状，否则TF无法推断，后续计算必报错
        output.set_shape(inputs.shape)
        return output
    
    # 可选：也可以重写这个方法指定输出形状，效果和set_shape一致
    def compute_output_shape(self, input_shape):
        return input_shape

步骤3：测试层功能

# 构建简单模型测试
input_layer = tf.keras.Input(shape=(3,))
custom_layer = MyCustomLayer()(input_layer)
model = tf.keras.Model(inputs=input_layer, outputs=custom_layer)

test_input = np.array([[1, 2, 3], [4, 5, 6]])
print(model.predict(test_input))
# 输出应为：[[ 2.  5. 10.] [17. 26. 37.]]

场景2：需要支持自动微分（用于训练） #

上面的方法无法自动计算梯度，如果要把层用在训练模型里，就得用tf.custom_gradient手动定义梯度：

@tf.custom_gradient
def my_func_with_gradient(input_tensor):
    # 前向传播：将TF张量转为numpy数组处理
    input_np = input_tensor.numpy()
    result_np = input_np ** 2 + 1
    result = tf.convert_to_tensor(result_np, dtype=tf.float32)
    
    # 定义梯度函数：dy是上游传来的梯度，这里计算当前层的梯度
    def gradient_fn(dy):
        # 示例函数y=x²+1的梯度是2x，所以返回 dy * 2x
        return dy * 2 * input_tensor
    
    return result, gradient_fn

class TrainableMyLayer(tf.keras.layers.Layer):
    def call(self, inputs):
        return my_func_with_gradient(inputs)

这个层可以正常参与训练，TF能正确计算梯度。

为什么之前的tf.py_func会报错？ #

tf.py_func在TF2.x已被标记为过时，它最大的问题是无法自动推断输出张量的形状，导致后续层计算时，形状里出现None值。当代码尝试对这个形状的整数元素做%=操作时，就会触发int与NoneType不兼容的错误。而tf.numpy_function通过手动指定形状，就能避免这个问题。

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