如何为TensorFlow Lite中封装scipy.lfilter的自定义算子设置正确名称

你好，我来帮你解决这个问题。你现在遇到的核心问题是：用tf.numpy_function封装scipy的lfilter后，导出TFLite时算子被标记为通用的PyFunc，而不是你想要的Lfilter。这是因为tf.numpy_function本质是TF提供的Python/numpy代码桥接工具，TFLite会把所有通过它实现的函数统一识别为PyFunc类型，你传入的name="Lfilter"只是TF计算图里的节点名称，不是算子的类型名，所以不会被TFLite当作自定义算子名。

下面给你两种解决方案，按需选择：

方案一：规范自定义算子（适合长期维护的项目） #

这种方法需要你放弃tf.numpy_function，改用TF官方的自定义算子注册机制，这样导出TFLite时算子名会被直接识别为你想要的名称，同时能更好地兼容Keras的序列化逻辑。

修改后的完整核心代码 #

import tensorflow as tf
import keras
from scipy import signal
import numpy as np

LFILTER_COEFF_DTYPE = np.float32
LFILTER_DATA_DTYPE = np.float32

# 1. 定义带梯度的自定义lfilter实现（如果不需要训练，梯度可以返回全零）
@tf.custom_gradient
def custom_lfilter(b, a, x):
    # 前向计算复用scipy的lfilter
    y_np = signal.lfilter(b.numpy(), a.numpy(), x.numpy())
    y = tf.convert_to_tensor(y_np, dtype=LFILTER_DATA_DTYPE)
    y.set_shape(x.shape)

    # 反向传播：如果不需要训练，直接返回全零张量
    def grad(dy):
        return tf.zeros_like(b), tf.zeros_like(a), tf.zeros_like(x)
    
    return y, grad

# 2. 包装为TF可追踪的函数
@tf.function(input_signature=[
    tf.TensorSpec(shape=[None], dtype=LFILTER_COEFF_DTYPE),
    tf.TensorSpec(shape=[None], dtype=LFILTER_COEFF_DTYPE),
    tf.TensorSpec(shape=[None, None], dtype=LFILTER_DATA_DTYPE),
])
def tf_lfilter(b, a, x):
    return custom_lfilter(b, a, x)

# 3. 完善Keras层的序列化逻辑
@keras.saving.register_keras_serializable()
class Lfilter(keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.b = self.add_weight(
            name="b",
            shape=(3,),
            initializer="ones",
            trainable=True,
            dtype=LFILTER_COEFF_DTYPE
        )
        self.a = self.add_weight(
            name="a",
            shape=(3,),
            initializer="ones",
            trainable=True,
            dtype=LFILTER_COEFF_DTYPE
        )

    def get_config(self):
        config = super().get_config()
        config.update({
            "b": self.b.numpy(),
            "a": self.a.numpy(),
        })
        return config

    @classmethod
    def from_config(cls, config):
        # 显式实现权重加载，确保名称和值正确恢复
        layer = cls(**config)
        layer.b.assign(config["b"])
        layer.a.assign(config["a"])
        return layer

    def call(self, x, training=None):
        y = tf_lfilter(self.b, self.a, x)
        y.set_shape(x.shape)
        return y

# 4. 导出TFLite的函数
def convert_to_tflite(model, input_shape, name=None):
    if name is not None:
        model.name = name
    model.build(input_shape)
    concrete_func = tf.function(model.call).get_concrete_function(
        tf.TensorSpec(input_shape, LFILTER_DATA_DTYPE)
    )
    converter = tf.lite.TFLiteConverter.from_concrete_functions(
        [concrete_func], model
    )
    converter.allow_custom_ops = True
    # 声明支持的自定义算子
    converter.target_spec.supported_custom_ops = ["Lfilter"]
    
    tflite_model = converter.convert()
    with open(f"{model.name}.tflite", "wb") as f:
        f.write(tflite_model)

方案二：快速修改导出过程（适合快速验证） #

如果不想修改TF侧的算子实现，只想快速把PyFunc替换成自定义算子名Lfilter，可以通过在导出TFLite时添加MLIR Pass来修改算子类型，不需要写C++代码。

修改后的导出函数（其余代码保留你的原有实现） #

def convert_to_tflite(model, input_shape, name=None):
    if name is not None:
        model.name = name
    tf_callable = tf.function(
        model.call,
        autograph=False,
        input_signature=[tf.TensorSpec(input_shape, LFILTER_DATA_DTYPE)],
    )
    tf_concrete_function = tf_callable.get_concrete_function()
    
    # 定义MLIR Pass：把PyFunc替换为自定义算子Lfilter
    def pyfunc_to_custom_op(mlir_module):
        import mlir
        from mlir.dialects import tf
        from mlir.ir import StringAttr
        
        for op in mlir_module.body.operations:
            if isinstance(op, tf.PyFuncOp):
                # 修改算子类型名为Lfilter
                op.name = StringAttr.get("Lfilter")
                # 移除PyFunc特有的属性
                op.attributes.pop("token", None)
                op.attributes.pop("device", None)
        return mlir_module
    
    converter = tf.lite.TFLiteConverter.from_concrete_functions(
        [tf_concrete_function], tf_callable
    )
    converter.allow_custom_ops = True
    # 启用MLIR转换器并添加自定义Pass
    converter.experimental_enable_mlir_converter = True
    converter.experimental_mlir_passes = [pyfunc_to_custom_op]
    
    tflite_model = converter.convert()
    with open(f"{model.name}.tflite", "wb") as f:
        f.write(tflite_model)

额外修复权重名称序列化问题 #

在你的Lfilter类中添加以下方法，确保权重的名称和值被正确序列化：

@classmethod
def from_config(cls, config):
    layer = cls(**config)
    layer.b.assign(config["b"])
    layer.a.assign(config["a"])
    return layer

用以上任意一种方法导出后，你再用Netron打开TFLite模型，就能看到算子名称变成Lfilter，权重b和a的名称也会正确显示。

备注：内容来源于stack exchange，提问作者EDL