MSE loss in tensorflow 2.0 mistakes y_true for a reduction key

Question

I am using a really simple neural network with the latest version of tensorflow 2.0, on a jupyter notebook running python 3.7.0. The NN has Xip, a float as output, which I use as a parameter in my function MainGaussian_1_U, which approximates an image. When I try to compute the loss using MeanSquareError between the real image img and the approximation mk, I am given an error in which the loss function seems to take img as a reduction key. After searches, I still have no idea what this key is supposed to be, and can't find a way to debug my code:

model = tf.keras.models.Sequential()

# Add the layers
model.add(tf.keras.layers.Dense(64, activation="relu"))
model.add(tf.keras.layers.Dense(32, activation="relu"))
model.add(tf.keras.layers.Dense(1, activation="relu"))

# The loss method
loss_object = tf.keras.losses.MeanSquaredError()
# The optimize
optimizer = tf.keras.optimizers.Adam()
# This metrics is used to track the progress of the training loss during the training
train_loss = tf.keras.metrics.Mean(name='train_loss')

def train_step(Data, img):
    MainGaussian_init(Data)
    for _ in range (5):
        with tf.GradientTape() as tape:
            Xip= model( (sizeh**-2 * np.ones((sizeh, sizeh))).reshape(-1, 49))
            MainGaussian_1_U ()
            print ("img=", img)
            loss= tf.keras.losses.MeanSquaredError(img, mk)
            print ("loss=", loss)
        gradients = tape.gradient(loss, model.trainable_variables)
        print (gradients)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
        train_loss(loss)

train_step (TestFile, TestFile[4])

The error given is:

c:\program files\python37\lib\site-packages\tensorflow_core\python\ops\losses\loss_reduction.py:67: FutureWarning: elementwise comparison failed; returning scalar instead, but in the future will perform elementwise comparison
  if key not in cls.all():

...

ValueError: Invalid Reduction Key [[21.05224609 20.79420471 34.9659729  ... 48.09233093 68.83874512
  83.10766602]
 [20.93516541 17.0511322  39.00476074 ... 56.74258423 47.75274658
  98.57067871]
 [38.18562317 22.70791626 24.37176514 ... 64.9606781  47.65338135
  67.61506653]
 ...
 [85.76565552 79.45443726 73.64129639 ... 73.66456604 47.06422424
  49.44664001]
 [87.14616394 82.38183594 77.00856018 ... 66.21652222 71.32862854
  58.39285278]
 [36.74142456 37.27145386 34.52891541 ... 29.58699036 37.37667847
  30.25990295]].

This is my first question here on Stack Overflow: please let me know if I can make it any clearer!

Answer 1

You correctly create the "loss object", but never use it. Instead, your code tries to create a new "loss object" with the images as parameters (which doesn't work). Instead, you want to put the images into the already-created loss object. You just have to change this line

loss= tf.keras.losses.MeanSquaredError(img, mk)

to

loss= loss_object(img, mk)