Slicing-based assignment in Keras / Tensorflow?

Question

Behold my Keras custom loss function:

def custom_loss(y_true, y_pred):
    sqerr = (y_true - y_pred)**2
    sqerr[:,4:-1:7] = sqerr[:,4:-1:7] * ((y_true[:,2:-1:7]-y_true[:,3:-1:7])/y_true[:,2:-1:7])**2
    return sqerr.mean()

But 'sqerr' is not a numpy array so this code results in the error

TypeError: 'Tensor' object does not support item assignment

So I read the post "How to do slice assignment in Tensorflow", including jdehesa's answer and the GitHub page on this discussion. And so here's what I've got now...

def custom_loss(y_true, y_pred):
    sqerr = K.square(y_true-y_pred)

    sqerr = tf.Variable(  sqerr , validate_shape=False )
    with tf.control_dependencies([sqerr[:,4:-1:7].assign( sqerr[:,4:-1:7] * ((y_true[:,2:-1:7]-y_true[:,3:-1:7])/y_true[:,2:-1:7])**2 )]):
        sqerr = tf.identity(sqerr)

    return K.mean(sqerr)

...but apparently I'm botching it when it actually comes time to use it:

Traceback (most recent call last): File "my_awesome_nn.py", line 119, in setup_model

model.compile(loss=custom_loss, optimizer=opt)  

File "/opt/anaconda/envs/py35/lib/python3.5/site-packages/keras/engine/training.py", line 850, in compile

sample_weight, mask)   

File "/opt/anaconda/envs/py35/lib/python3.5/site-packages/keras/engine/training.py", line 465, in weighted

score_array = K.mean(score_array, axis=list(range(weight_ndim, ndim))) 

TypeError: 'NoneType' object cannot be interpreted as an integer

What's happening is, the TF slicing is only allowed to be applied to Variables, not general Tensors, so I'm casting to a Variable. But when I cast to Variable, it wants to know the shape, but the shape is 'dynamically defined' at that point (i.e. first element is '?'). So setting validate_shape=False lets me actually define a variable, but this destroys the dimension info that Keras wants later. Observe:

def custom_loss(y_true, y_pred):  
        sqerr = K.square(y_true-y_pred)
        print("K.ndim(sqerr) #1 = ",K.ndim(sqerr))
        sqerr = tf.Variable(  sqerr , validate_shape=False )
        print("K.ndim(sqerr) #2 = ",K.ndim(sqerr))
        with tf.control_dependencies([sqerr[:,4:-1:7].assign( sqerr[:,4:-1:7] * ((y_true[:,2:-1:7]-y_true[:,3:-1:7])/y_true[:,2:-1:7])**2 )]):
            sqerr = tf.identity(sqerr)

        return K.mean(sqerr)

...results in the output

K.ndim(sqerr) #1 = 2

K.ndim(sqerr) #2 = None

Thus later, when the Keras training.py code is saying 'ndim = K.ndim(score_array)' it ends up with None, and thus the NoneType error.

Can anyone shed light on how to do what I need? Seems I can't slice without converting to a Variable, can't define a Variable for a dynamically-shaped tensor that will preserve the dynamic shape.

(This is a code that completely works if I just omit the "middle 3 lines" above and have my custom loss be regular MSE)

Answer 1

I think slice assignment can be avoided for this custom loss. If you want to adjust the loss value for sqerr[:, 4:-1:7], you can subtract the original value from the total loss, and then add back the adjusted loss value.

def custom_loss_keras(y_true, y_pred):
    # first sum up the squared error column-wise
    sqerr = K.square(y_true - y_pred)
    loss = K.sum(sqerr, axis=-1)

    # subtract the loss for the sliced part
    loss -= K.sum(sqerr[:, 4:-1:7], axis=-1)

    # add back the adjusted loss for the sliced part
    denominator = K.maximum(y_true[:, 2:-1:7], K.epsilon())  # handle zero-division
    numerator = y_true[:, 2:-1:7] - y_true[:, 3:-1:7]
    loss += K.sum(sqerr[:, 4:-1:7] * K.square(numerator / denominator), axis=-1)

    # take average
    ncols = K.int_shape(y_pred)[-1]
    loss /= ncols
    return K.mean(loss)

You can verify this function by comparing it to your original numpy version:

def custom_loss_numpy(y_true, y_pred):
    sqerr = (y_true - y_pred)**2
    sqerr[:,4:-1:7] = sqerr[:,4:-1:7] * ((y_true[:,2:-1:7]-y_true[:,3:-1:7])/y_true[:,2:-1:7])**2
    return sqerr.mean()

y_true = np.random.rand(50, 1000)
y_pred = np.random.rand(50, 1000)

print(custom_loss_numpy(y_true, y_pred))
889.992075384

print(K.eval(custom_loss_keras(K.variable(y_true), K.variable(y_pred))))
889.992