Пользовательский слой 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))

, таким образом, мой вопрос, как указано выше, касается этого процесса умножения.Я мог бы задать этот вопрос плохо, если да, то мои извинения.

Заранее большое спасибо:)