Tensorflow Keras create mask for labels

Question

I'm trying to create a model where my label data consists of floats and some of them are NaN. I can't just drop these labels from the dataset because it will be a recurrent network. Therefore I want to fill the NaN data with zero and create a mask for label NaN's so that it won't contribute to the loss function. Here is a toy example where I tried to create a mask on the output so that the NaN values are ignored.

import tensorflow as tf
import numpy as np
import tensorflow.keras.backend as K
tf.config.experimental_run_functions_eagerly(True)
print(tf.executing_eagerly())
print(tf.__version__)

# y_true = [np.nan, 0]
# break this up into a mask
# and fill np.nan with 0
mask = np.array([0., 1.])   # use 0's and 1's tensorflow doesn't support bool
y_true = np.array([0., 0.])

# create random model prediction
y_pred = np.array([1., 2.])

# define inputs for model
inputs = {}
inputs['x'] = tf.keras.layers.Input(shape=(None,), name='x')
inputs['mask'] = tf.keras.layers.Input(shape=(None,), name='mask')

class OutputMask(tf.keras.layers.Layer):
  # create a masking layer class
  # assumes shape [:, 0] = input
  # [:,1] = mask
  def call(self, inputs):
    return inputs[:, 0:1]
    
  def compute_mask(self, inputs, mask=None):
    self.mask = tf.cast(inputs[:,-1:], 'bool')
    return mask

# toy model where output = input with mask
out = tf.keras.layers.concatenate([inputs['x'], inputs['mask']])
mask_layer = OutputMask()
output = mask_layer(out)
model = tf.keras.Model(inputs=inputs, outputs=output)

def masked_loss(mask_layer):
  # defines custom loss function that supports masking
  # needs masking layer applied to it
  mask = mask_layer.mask
  mse = tf.keras.losses.MeanSquaredError()
  def loss(y_true, y_pred):
    y_true = tf.boolean_mask(y_true, mask)
    y_pred = tf.boolean_mask(y_pred, mask)
    return mse(y_true, y_pred)
  return loss

model.compile(optimizer="Adam", loss=masked_loss(mask_layer))
model.evaluate(x={'x':y_pred, 'mask':mask},y=y_true)

# returns zero instead of expected result 4.

I want my loss function to be (0-2)^2 = 4. The model actually returns 0. for a loss but for some reason If I run compile twice (using Jupyter notebook) the model returns 4 as expected.

Therefore I have two questions:

Why am I having to compile the model twice before it works
Is there a better way to achieve what I want?

Tensorflow Keras create mask for labels

0 Answers