在Keras层中执行傅立叶卷积时没有可训练的参数

时间:2019-11-21 23:08:13

标签: tensorflow keras convolution

我正在尝试使用tf.keras实现傅立叶卷积神经网络,其中将输入和内核转换到频域,执行逐元素乘法,然后对输出进行逆变换和裁剪。模型摘要显示,即使我使用self.add_weight声明了FConv2D层中的内核,也没有可训练的参数。应该有(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不能在自定义层中进行这些复杂的操作吗?

1 个答案:

答案 0 :(得分:0)

问题在于我如何重命名变量。我将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。