custom loss function in Keras with masking array as input

user_1731970 0
Viewed 1767

I am trying to train an Autoencoder with a custom loss function shown below. The input, missing_matrix, is an n x m array of 1s and 0s corresponding to the n x m features array. I need to do an element by element multiplication of the missing_array with y_pred, which should be a reconstruction of the input features so that I can mask those that get multiplied by 0 to neglect their contribution in the cost function. I have never written a custom loss function before, the one below doesn't work at all. I have tried to search for similar custom cost functions but haven't been able to find one that brings in some input array like this. I would appreciate the help or a point in the right direction.

def custom_loss(missing_array):
    
    def missing_mse(y_true, y_pred):
        mse = MeanSquaredError()
        y_pred_masked = tf.math.multiply(y_pred, missing_array)
        return mse(y_true = y_true, y_pred = y_pred_masked)

    return missing_mse

Edit: Got a bit further

from keras.losses import MeanSquaredError
import tensorflow as tf

def custom_loss(missing_matrix):
    
    def missing_mse(y_true, y_pred):
        mse = MeanSquaredError()
        y_pred_masked = tf.math.multiply(y_pred, tf.convert_to_tensor(missing_matrix, dtype=tf.float32))
        return mse(y_true = y_true, y_pred = y_pred_masked)

    return missing_mse

with error

InvalidArgumentError:  Incompatible shapes: [64,1455] vs. [13580,1455]
     [[node gradient_tape/missing_mse/BroadcastGradientArgs (defined at <ipython-input-454-b60d74568bf2>:64) ]] [Op:__inference_train_function_25950]

Function call stack:
train_function

The 64 makes me think it is the batch. Likely I need to take batches of 64 of my missing matrix?

Edit2: This is fun! So I verified that the custom loss function will train if I do something like

def train(self, model, X_train):
        """
        Model training
        """
        #model.fit(X_train, X_train, epochs = 10, batch_size = 64, validation_split = 0.10)
        for batch_idx in range(0, len(X_train), 70):
            self.batch_start = batch_idx
            self.batch_end = batch_idx + 70
            model.train_on_batch(X_train[self.batch_start:self.batch_end,:], X_train[self.batch_start:self.batch_end,:])
        return model

and modify my custom loss

def custom_loss2(self, missing_matrix):
    
        def missing_mse(y_true, y_pred):
            mse = MeanSquaredError()
            y_pred_masked = tf.math.multiply(y_pred, tf.convert_to_tensor(missing_matrix[self.batch_start:self.batch_end,:], dtype=tf.float32))
            return mse(y_true = y_true[self.batch_start:self.batch_end,:], y_pred = y_pred_masked[self.batch_start:self.batch_end,:])

        return missing_mse

So now how can I get epochs and print out validation loss etc...? Or rather what is the better way to do this? Thats it for me tonight. Goodnight!

keras tensorflow deep-learning machine-learning python
1 Answers

The problem is that y_true and y_pred are in batches while the mask is passed one-shot. One simple solution to automatically split your data into equal batches is using model.add_loss().

Below I reproduced a dummy example with an autoencoder and a custom masking loss. The mask is passed as model input and this is the simple trick you need to make it work.

Use also mean_squared_error instead of MeanSquaredError.

def missing_mse(y_true, y_pred, missing_array):
    
    y_pred_masked = tf.math.multiply(y_pred, missing_array)
    
    mse = tf.keras.losses.mean_squared_error(y_true = y_true, y_pred = y_pred_masked)
    
    return mse

inp = Input((100,))
x = Dense(20)(inp)
out = Dense(100)(x)

inp_mask = Input((100,))

model = Model([inp, inp_mask], out)
model.add_loss(missing_mse(inp, out, inp_mask))
model.compile('adam', loss=None)

X = np.random.uniform(0,1, (300,100))
mask = np.random.randint(0,2, (300,100))

model.fit(x=[X,mask], y=None, epochs=10, batch_size=64)

# at inference time you can remove the mask input in this way
new_model = Model(model.input[0], model.output)
new_model.predict(X).shape
user_103750490
Related