Create a weighted MSE loss function in Tensorflow

I want to underline that you have 2 possibilities according to your problem:

[1] If the weights are equal for all your samples:

You can build a loss wrapper. Here a dummy example:

n_sample = 200
X = np.random.uniform(0,1, (n_sample,10))
y = np.random.uniform(0,1, (n_sample,100))
W = np.random.uniform(0,1, (100,)).astype('float32')

def custom_loss_wrapper(weights):
    def loss(true, pred):
        sum_weights = tf.reduce_sum(weights) * tf.cast(tf.shape(pred)[0], tf.float32)
        resid = tf.sqrt(tf.reduce_sum(weights * tf.square(true - pred)))
        return resid/sum_weights
    return loss

inp = Input((10,))
x = Dense(256)(inp)
pred = Dense(100)(x)

model = Model(inp, pred)
model.compile('adam', loss=custom_loss_wrapper(W))

model.fit(X, y, epochs=3)

[2] If the weights are different between samples:

You should build your model usind add_loss in order to dinamically take into account the weights for each sample. Here a dummy example:

n_sample = 200
X = np.random.uniform(0,1, (n_sample,10))
y = np.random.uniform(0,1, (n_sample,100))
W = np.random.uniform(0,1, (n_sample,100))

def custom_loss(true, pred, weights):
    sum_weights = tf.reduce_sum(weights)
    resid = tf.sqrt(tf.reduce_sum(weights * tf.square(true - pred)))
    return resid/sum_weights

inp = Input((10,))
true = Input((100,))
weights = Input((100,))
x = Dense(256)(inp)
pred = Dense(100)(x)

model = Model([inp,true,weights], pred)
model.add_loss(custom_loss(true, pred, weights))
model.compile('adam', loss=None)

model.fit([X,y,W], y=None, epochs=3)

When using add_loss you should pass all the tensors involved in the loss as input layer and pass them inside the loss for the computation.

At inference time you can compute predictions as always, simply removing the true and weights as input:

final_model = Model(model.input[0], model.output)
final_model.predict(X)

You can implement custom weighted mse in the following way

import numpy as np 
from tensorflow.keras import backend as K 

def custom_mse(class_weights):
    def weighted_mse(gt, pred):
        # Formula: 
        # w_1*(y_1-y'_1)^2 + ... + w_100*(y_100-y'_100)^2 / sum(weights)
        return K.sum(class_weights * K.square(gt - pred)) / K.sum(class_weights)
    return weighted_mse

y_true  = np.array([[0., 1., 1, 0.], [0., 0., 1., 1.]])
y_pred  = np.array([[0., 1, 0., 1.], [1., 0., 1., 1.]])
weights = np.array([0.25, 0.50, 1., 0.75])

print(y_true.shape, y_pred.shape, weights.shape)
(2, 4) (2, 4) (4,)

loss = custom_mse(class_weights=weights)
loss(y_true, y_pred).numpy()
0.8

Using it with model compilation.

model.compile(loss=custom_mse(weights))

This will compute mse with the provided weighted matrices. However, in your question, you quote sqrt..., from which I presume you meant root mse (rmse). To do that you can use K.sqrt(K.sum(...)) / K.sum(...) in the custom function of custom_mse.

FYI, you may also interest to look at class_weights and sample_weights during Model. fit. From source:

• class_weight: Optional dictionary mapping class indices (integers) to a weight (float) value, used for weighting the loss function (during training only). This can be useful to tell the model to "pay more attention" to samples from an under-represented class.

• sample_weight: Optional Numpy array of weights for the training samples, used for weighting the loss function (during training only). You can either pass a flat (1D) Numpy array with the same length as the input samples (1:1 mapping between weights and samples), or in the case of temporal data, you can pass a 2D array with shape (samples, sequence_length), to apply a different weight to every timestep of every sample. This argument is not supported when x is a dataset, generator, or keras.utils.Sequence instance, instead provides the sample_weights as the third element of x.

And also loss_weights in Model.compile, from source

loss_weights: Optional list or dictionary specifying scalar coefficients (Python floats) to weight the loss contributions of different model outputs. The loss value that will be minimized by the model will then be the weighted sum of all individual losses, weighted by the loss_weights coefficients. If a list, it is expected to have a 1:1 mapping to the model's outputs. If a dict, it is expected to map output names (strings) to scalar coefficients.

A class version of the weighted mean squared error loss function.

class WeightedMSE(object):
    def __init__(self):
        pass

    def __call__(self, y_true, y_pred, weights):
        sum_weights = tf.reduce_sum(weights)
        resid = tf.reduce_sum(weights * tf.square(y_true - y_pred))
        return resid / sum_weights