python - Keras自定义层：如何将层的输入与自定义层的定义参数相乘？

Question

我正在尝试在 Keras 中创建一个自定义层，它接受输入并将其与作为参数传递给层的张量相乘。我还做了各种操作，例如 fft 和 ifft（以及它们各自的班次），这些操作与我认为的这个问题无关，因为我已经测试了它们的正确性，所以我需要定义一个自定义层。

我已经尝试逐行调试我的代码，它似乎可以工作，但是当我在模型中对其进行测试时，该层并没有给出预期的结果。我还尝试获取变量的形状并尝试通过挤压参数进行调整，但无济于事。

inputs1 = Input(shape=(256, 256, 1), name='inputs')
inputs2 = Input(shape=(256, 256, 1), name='k_space_real')
inputs3 = Input(shape=(256, 256, 1), name='k_space_imag')

concat1 = concatenate([inputs2, inputs3], axis=3)

concat2 = concatenate([inputs1, concat1], axis=3)

outputs = CustomLayer(arg)(concat2)

self.model = Model(inputs=[inputs1, inputs2, inputs3], outputs=[outputs])
#################################   
class CustomLayer(layers.Layer):
    def __init__(self, mask, **kwargs):
        super(UpdatedDCLayer, self).__init__(**kwargs)
        self.mask = mask

    def call(self, inputs):
        a = inputs[:, :, :, 0] 
        input_k_space_real = inputs[:, :, :, 1] 
        input_k_space_imag = inputs[:, :, :, 2] 
        input_k_space = tf.complex(input_k_space_real, input_k_space_imag) 
        input_k_space = tf.cast(input_k_space, dtype=tf.complex128)

        mask_bool = tf.cast(self.mask, dtype=tf.bool) 
        inv_mask = tf.cast(tf.math.logical_not(mask_bool), dtype=tf.complex128) 

        intermediary = tf.cast(a, dtype=tf.complex128)
        intermediary_k_space = tf.spectral.fft2d(intermediary)
        intermediary_k_space = tf.roll(intermediary_k_space, shift=[128, 128], axis=[0, 1]) 

        output_intermediary = tf.math.multiply(intermediary_k_space, tf.squeeze(inv_mask))
        input_dc_k_space = tf.math.add(output_intermediary, input_k_space)
        input_dc_k_space = tf.roll(input_dc_k_space, shift=[-128, -128], axis=[0, 1])
        output_dc = tf.spectral.ifft2d(input_dc_k_space)
        output_intermediary = tf.roll(output_intermediary, shift=[-128, -128], axis=[0, 1]) 
        output_intermediary = tf.spectral.ifft2d(output_intermediary)
        output_dc = tf.expand_dims(output_dc, -1) 
        output_intermediary = tf.expand_dims(output_intermediary, -1) 
        return [tf.cast(tf.math.abs(output_dc), dtype=tf.float32),
            tf.cast(tf.math.abs(output_intermediary), dtype=tf.float32)]



    def compute_output_shape(self, input_shape):
        shape = list(input_shape)
        shape[-1] = 1
        return [tuple(shape), tuple(shape)]

该层将在模型期间多次使用，并且 input1 依赖于之前的 CNN，其中 input2 和 3 是常量。所有输入具有相同的形状 [sample_size, 256, 256, 1]。

这个自定义层给出了一个结果，但不是预期的。问题似乎源于这一行：

output_intermediary = tf.math.multiply(intermediary_k_space, tf.squeeze(inv_mask))

因此，如上所述，我的问题是关于这个乘法过程。我可能问得不好，如果是这样，我很抱歉。

提前非常感谢:)

标签： pythontensorflowkeras