Custom Loss Function That Can Accept Missing Values in Tensorflow

Question

I am creating a custom semi-supervised learning model using the functional API from Tensorflow 2. Suppose I have a y_true of a single sample as:

y_true = np.array[[0.1, np.nan],[np.nan,0.2]]

Where np.nan represent unobserved data. I want to create a custom Mean Square Error loss function to take the loss of the observed data only, ie:

def customloss(y_true, y_pred):
    observed_loss_matrix = tf.zeros_like(y_true) #create [0 0] matrix
    boolean_true = tf.math.is_nan(y_true) #get boolean matrix
    observed_loss_matrix[~boolean_true] = 1 
    y_true[boolean_true] = 0 #I'm afraid np.nan might mess up with the gradient computation so I change it to 0
    error = (y_true - y_pred)**2
    observed_error = tf.math.multiply(error,observed_loss_matrix) #only takes
    mse = tf.math.reduce_sum(observed_error)
    return mse

I wonder if it's the right way to do it. I highly appreciate any help.

Answer 1

Although it work:

def mse_mv(y_true, y_pred):
     per_instance = tf.where(tf.is_nan(y_true),
                             tf.zeros_like(y_true),
                             tf.square(tf.subtract(y_pred, y_true)))
     return tf.reduce_mean(per_instance, axis=0)

Or solution with keras

import keras.backend as K

def custom_error_function(y_true, y_pred):
    bool_finite = T.is_finite(y_true)
    return K.mean(K.square(T.boolean_mask(y_pred, bool_finite) - T.boolean_mask(y_true, bool_finite)), axis=-1)

Or:

from keras.models import Model
from keras.layers import Input, Dense
import keras.backend as K

inp = Input(shape=(3,))

w1 = Input(shape=(1,))
w2 = Input(shape=(1,))

out1 = Dense(1)(inp)
out2 = Dense(1)(inp)

def weighted_loss(weight):

    def loss(y_true,y_pred):
        return K.mean(K.square(y_pred - y_true) * weight , axis=-1)
    return loss

model = Model(inputs=inp, outputs=[out1,out2])

modelw = Model(inputs=[inp,w1,w2],outputs=[out1,out2])
modelw.compile(optimizer='rmsprop',
              loss=[weighted_loss(w1),weighted_loss(w2)])

Train

import numpy as np
inp_v = np.array([[0,0,0],[1,1,1],[0,1,1],[1,0,0],[1,1,0]])

out1_v = np.array([0,0,0,0,0])
out2_v = np.array([1,np.nan,np.nan,1,1])

w1_v = np.array(~np.isnan(out1_v),dtype=np.int)
w2_v = np.array(~np.isnan(out2_v),dtype=np.int)

#Replace nan by 0 or anything just in case
out1_v[np.isnan(out1_v)] = 10000
out2_v[np.isnan(out2_v)] = 10000

modelw.fit([inp_v,w1_v,w2_v],[out1_v,out2_v],epochs=1000,verbose=False)