在 Keras 层中实现傅里叶卷积时没有可训练的参数

Question

我正在尝试使用 tf.keras 实现傅立叶卷积神经网络，其中输入和内核被转换到频域，执行元素乘法，然后输出被逆变换和裁剪。模型摘要显示在我的 FConv2D 层中没有可训练的内核参数，即使我使用 self.add_weight 声明它们。应该有(3*3*in_channels*no_of_kernels)个参数。

class FConv2D(tf.keras.layers.Layer):
    def __init__(self, no_of_kernels, **kwargs):
        self.no_of_kernels = no_of_kernels
        super(FConv2D, self).__init__(**kwargs)

    def build(self, input_shape):
        self.kernel_shape = (3, 3 , int(input_shape[3]), self.no_of_kernels)
        print(input_shape, self.kernel_shape)
        self.kernel = self.add_weight(shape=(3,3, input_shape[-1], self.no_of_kernels),
                             initializer='random_normal',
                             trainable=True)
        self.bias = self.add_weight(shape=(self.no_of_kernels,),
                             initializer='random_normal',
                             trainable=True)
        super(FConv2D, self).build(input_shape)


    def call(self, x):
        crop_size = self.kernel.get_shape().as_list()[0] // 2
        shape = x.get_shape().as_list()[1] + self.kernel.get_shape().as_list()[0] - 1
        x = tf.transpose(x, perm=[0,3,1,2])
        self.kernel = tf.transpose(self.kernel, perm=[3,2,0,1])
        x = tf.signal.rfft2d(x, [shape, shape])
        self.kernel = tf.signal.rfft2d(self.kernel, [shape, shape])
        x = tf.einsum('imkl,jmkl->ijkl', x, self.kernel)
        x = tf.signal.irfft2d(x, [shape, shape])
        x = tf.transpose(x, perm=[0,2,3,1])
        x = tf.nn.bias_add(x, self.bias)[:,crop_size:-1*crop_size, crop_size:-1*crop_size, :]
        x = tf.nn.elu(x)
        return x      

当我构建模型时，它只显示偏差项的可训练参数，而不是内核。

m = tf.keras.models.Sequential()
m.add(FConv2D(32, input_shape=(32,32,3)))
m.summary()

输出：

Model: "sequential_37"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
f_conv2d_88 (FConv2D)        (None, 32, 32, 32)        32        
=================================================================
Total params: 32
Trainable params: 32
Non-trainable params: 0

问题似乎出在call(self, x)，因为如果我将傅里叶卷积运算替换为对tf.nn.conv2d 的调用，则会列出预期的参数数量（3*3*3*32+32=896）。

我已通过消除偏差项并调用model.fit 确认参数不可训练，因为没有可训练的参数，因此无法运行。

我错过了什么？ Keras 不能在自定义层中进行这些复杂的操作吗？

Answer 1

问题在于我如何重命名变量。我将self.kernel 设置为转换后的变量，因此当我对转换中的所有操作使用不同的变量W 时，参数会按预期列出。

作为参考，下面列出的代码是 Conv2D 层的替代代码。

class FConv2D(tf.keras.layers.Layer):
    def __init__(self, no_of_kernels, kernel_shape, **kwargs):
        self.no_of_kernels = no_of_kernels
        self.kernel_shape = kernel_shape
        super(FConv2D, self).__init__(**kwargs)

    def build(self, input_shape):
        self.kernel = self.add_weight(shape=self.kernel_shape + (input_shape[-1], self.no_of_kernels),
                             initializer='random_normal',
                             trainable=True)
        self.bias = self.add_weight(shape=(self.no_of_kernels,),
                             initializer='random_normal',
                             trainable=True)
        super(FConv2D, self).build(input_shape)


    def call(self, x):
        crop_size = self.kernel.get_shape().as_list()[0] // 2
        shape = x.get_shape().as_list()[1] + self.kernel.get_shape().as_list()[0] - 1
        X = tf.transpose(x, perm=[0,3,1,2])
        W = tf.transpose(self.kernel, perm=[3,2,0,1])
        X = tf.signal.rfft2d(X, [shape, shape])
        W = tf.signal.rfft2d(W, [shape, shape])
        X = tf.einsum('imkl,jmkl->ijkl', X, W)
        output = tf.signal.irfft2d(X, [shape, shape])
        output = tf.transpose(output, perm=[0,2,3,1])
        output = tf.nn.bias_add(output, self.bias)[:,crop_size:-1*crop_size, crop_size:-1*crop_size, :]
        return output   

使用 FConv2D(32, (3,3), input_shape=(32,32,3)) 替换 tf.keras.layers.Conv2D。