TensorFlow自定义1D卷积层变量无梯度警告原因及解决方法

我在TensorFlow中构建了一个自定义1D卷积层，功能验证正常，但加入Keras Sequential模型时，出现「自定义层中的变量不存在梯度」的警告。相关代码如下：

import tensorflow as tf
import numpy as np


class customC1DLayer(tf.keras.layers.Layer):
    def __init__(self, filter_size = 1, activation = None ,**kwargs):
        super(customC1DLayer, self).__init__(**kwargs)
        self.filter_size = filter_size
        self.activation = tf.keras.activations.get(activation)
    
    def build(self, input_shape):
        self.filter = self.add_weight('filter', shape=[self.filter_size, ], trainable=True, dtype=tf.float32)

        self.padding = tf.Variable(initial_value=tf.zeros(shape=[input_shape[-1] - self.filter_size, ], dtype=tf.float32), trainable=False)
        padded_filter = tf.concat([self.filter, self.padding], axis=0)
        col = tf.concat([padded_filter[:1], tf.zeros_like(padded_filter[1:])], axis=0)
        self.augmented_filter = tf.linalg.LinearOperatorToeplitz(padded_filter, col).to_dense()
     
    def call(self, inputs):
        outputs = tf.transpose(tf.matmul(self.augmented_filter, inputs, transpose_b=True))
        if self.activation is not None:
            outputs = self.activation(outputs)
        return outputs

代码说明：在build方法中初始化权重（如[a b c]），随后生成循环矩阵作为augmented_filter，例如[[a b c 0 0], [0 a b c 0], [0 0 a b c]]。已知此类问题常因使用不可微分函数导致，但此处仅使用了可微分的矩阵操作，对此感到困惑。

问题原因 #

核心问题出在**build方法中直接计算并赋值self.augmented_filter**：

• build仅在层初始化时执行一次，self.augmented_filter会被固化为静态张量，而非与self.filter绑定的动态计算节点。
• 后续call方法使用的是这个静态张量，梯度无法反向传播到原始的self.filter变量，导致TensorFlow判定该变量无梯度关联。
• 矩阵操作本身可微分，但提前固化计算结果直接切断了梯度流路径。

修复方法 #

需要将augmented_filter的计算逻辑从build移到call方法中，确保每次前向传播都基于当前self.filter动态生成循环矩阵，维持梯度流的连续性：

1. 移除build中关于augmented_filter的计算代码，仅保留可训练变量初始化。
2. 在call内动态生成填充后的滤波器、Toeplitz矩阵，让计算图与self.filter绑定。

修改后的代码：

import tensorflow as tf
import numpy as np


class customC1DLayer(tf.keras.layers.Layer):
    def __init__(self, filter_size = 1, activation = None ,**kwargs):
        super(customC1DLayer, self).__init__(**kwargs)
        self.filter_size = filter_size
        self.activation = tf.keras.activations.get(activation)
    
    def build(self, input_shape):
        # 仅初始化可训练权重，移除静态计算逻辑
        self.filter = self.add_weight('filter', shape=[self.filter_size, ], trainable=True, dtype=tf.float32)
        # 预计算padding长度，避免每次call重复解析input_shape
        self.pad_length = input_shape[-1] - self.filter_size
     
    def call(self, inputs):
        # 在call中动态生成augmented_filter，维持梯度流
        padding = tf.zeros(shape=[self.pad_length, ], dtype=tf.float32)
        padded_filter = tf.concat([self.filter, padding], axis=0)
        col = tf.concat([padded_filter[:1], tf.zeros_like(padded_filter[1:])], axis=0)
        augmented_filter = tf.linalg.LinearOperatorToeplitz(padded_filter, col).to_dense()
        
        outputs = tf.transpose(tf.matmul(augmented_filter, inputs, transpose_b=True))
        if self.activation is not None:
            outputs = self.activation(outputs)
        return outputs

额外说明 #

• 移除了self.padding这个不可训练变量，改为在call中生成临时张量，减少不必要的层变量。
• pad_length在build中预计算，避免每次call重复解析输入形状，提升运行效率。
• 动态计算确保augmented_filter始终与self.filter的当前值关联，梯度可正常反向传播到可训练权重。

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