Is there an established solution for accepting partially-NaN targets in Tensorflow?

Viewed 30

I have a dataset with multi-dimensional targets including many NaNs, and it is not possible to drop every sample containing NaNs in the target. Is the best solution to write a custom loss function that drops nan values? Such a function appears to work as expected, while utilizing the provided mean squared error loss results in an undefined loss value and gradients.

def masked_mse(y_actual,y_pred): 
    custom_loss = tf.keras.backend.square(y_actual-y_pred)
    custom_loss = tf.where(tf.math.is_nan(custom_loss), tf.zeros_like(custom_loss), custom_loss)
    return custom_loss

X, y = make_regression(n_samples=(n_samples + n_test_samples), n_features=n_features, n_targets=n_targets)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=n_test_samples)
nan_idx = np.ones(n_samples)
nan_idx[:nan_count_per_col] = np.nan

for n in range(n_targets):
  np.random.shuffle(nan_idx)
  y_train[:, n] *= nan_idx

model.compile(loss=masked_mse)
model.fit(X_train, y_train)
model.evaluate(X_test, y_test)
# Returns ~1E-2, a strong score

I'm surprised nothing came up through my Googling as I would imagine this is a fairly common need. Am I missing something?

0 Answers
Related