Custom loss function access tensor channels

Question

I have a 2-channel numpy array of shape (64, 64, 2) as input to my CNN. I want to build a customized loss function as described in https://www.tensorflow.org/guide/keras/train_and_evaluate :

def basic_loss_function(y_true, y_pred):
    return tf.math.reduce_mean(tf.abs(y_true - y_pred))

model.compile(optimizer=keras.optimizers.Adam(),
              loss=basic_loss_function)

model.fit(x_train, y_train, batch_size=64, epochs=3)

But I want to something more complicated that this basic one. What I need is to do an inverse DFT (ifft2d) and my y_pred and y_true are expected to be each of shape (64,64,2), with the 2 channels being the real and imaginary parts of a fft2. How can I access correctly the y_pred and y_true channels (which are some kind of keras/tensor layer I guess?) to rebuild a complex number in the form RealPart+1j*ImagPart (in numpy it would be y_pred[:,:,0] and y_pred[:,:,1] ) ?

--> In summary, does someone know exactly what kind of object is y_pred and y_true and how to access their channels/elements? (This is not easy to debug since would need to be run in a compiled CNN, so better know it beforehand)

Answer 1

y_true and y_pred are Tensors of shape (batchsize, ...[output shape]...). Your input has shape (64,64,2) but I'm not sure what your output looks like, if your output is indeed (64,64,2) then y_pred or y_true have shape (64,64,64,2) given your batchsize=64.

Dealing with Tensors is very much like numpy's syntax so You can use slice notation with tensors e.g., y_true[:,:,:,0] (note the added batch dimension).

Tensorflow has functions for computing DFT,FFT,..etc. See tf.signal and tf.signal.rfft2d

If your loss function involves operations on the input, not just the outputs y_true and y_pred, then you can use model.add_loss instead of model.compile(loss= basic_loss_function) as follows

x = Input(shape=(64,64,2))
y_true = Input(shape=...))
# your CNN layers
y_pred = Dense(128)(net)

model = Model(input=[x, y_true], output=y_pred)
model.add_loss(basic_loss_function(x, y_true, y_pred))

note that the labels (aka y_true) is now an input to the model.