Weighted mse custom loss function in keras - custom weights

Question

I'm working with sequence data, (one hot encoded sequences), and am looking for a way to write up a custom loss function that uses weights from a dictionary of values based on y_pred and y_true, and depends on those values while training (so I can't use constant weights when calling fit).

Basically, for each argmax index position in the sequence-matrix, I can retrieve a character. For each two characters i can retrieve a weight. Dictionaries for these are like below:

values = 
     {0: 'A',
      1: 'C',
      2: 'D',
      ...}
matrix = array
     ([[ 4, -1, -2, -2,  0, -1, -1,  0, -2, -1, -1, -1, -1, -2, -1,  1,
         0, -3, -2,  0, -2, -1,  0, -4],
       [-1,  5,  0, -2, -3,  1,  0, -2,  0, -3, -2,  2, -1, -3, -2, -1,
        -1, -3, -2, -3, -1,  0, -1, -4],
       [-2,  0,  6,  1, -3,  0,  0,  0,  1, -3, -3,  0, -2, -3, -2,  1,
         0, -4, -2, -3,  3,  0, -1, -4],
       ...]])

I want to do something like this:

y_true (n,155,20) ---K.argmax(.., axis=2)---> a:(n,155)

y_pred (n,155,20) ---K.argmax(.., axis=2)---> b:(n,155)

for i in range(n):
   for j in range(155):
      weights[i,j] = matrix[values[a[i,j]], values[b[i,j]]]

Imagine there is a way of getting the matrix values above via some other dict.

Then I wanna use my weights matrix like this:

def custom_loss_mse(y_true,y_pred):
    w = getWeights(y_true,y_pred)
    return K.mean(K.dot(w, K.square(y_pred-y_true)), axis=-1)

So far I have only found this question helpful, and it is not really similar.

This is easy, but keras makes it much harder for me to do because of the computation graph model. There should be some fast way to do this, but I am out of ideas.

I would appreciate any help since I am fairly new to numpy and keras.

Answer 1

You can use tf.gather to get values of a tensor at a specific index. For example, you could construct a with the following code:

index = K.argmax(y_true, axis=2)
a = tf.diag_part(tf.gather(y_true, index, axis=2))  # shape (n, 155)

I don't think it is possible to get values from a dict since all computations must be symbolic. But if you can use a matrix instead of a dict, you can convert the matrix to a tensor with K.constant and then you can use the same approach as above to get values at a specific index, depending on the input tensors.