Decorating a custom loss with tf.function changes the training results completely, both in keras model.fit method as well as custom training loop

Question

I was trying to implement class_weight myself as I saw some behavior in the usual way of doing it (class_weight in model.fit method) and it led me to a see a behavior that is very confusing and after thinking about it for quite some time I cannot understand what is happening.

The whole post is long so I give a brief description before giving the details for those interested

Brief Description

In brief, I define two custom loss functions that differ only in the tf.function decorator. The first one does not have the decorator while the second one does. Using keras model.fit, the training does not converge with the one that has the decorator. When I try the same with custom training loop the behavior is reversed.

I can obviously see which is the right behavior by running everything in eager mode and so I can see that the custom loop with the decorate function is the right one. Its results also match the usual keras model fit method with class_weight. All the other three methods give wrong results but I cannot understand why.

Detailed description

def customloss1(y_true,y_pred,sample_weight=None):
    weights=tf.constant([1.,1.,.1])[tf.newaxis,...]
    y_true_one_hot=tf.one_hot(tf.cast(y_true,tf.uint8),3)
    return tf.reduce_mean(tf.keras.losses.categorical_crossentropy(y_true_one_hot*weights,
                                                                   y_pred,
                                                                   from_logits=False))

@tf.function
def customloss2(y_true,y_pred,sample_weight=None):
    weights=tf.constant([1.,1.,.1])[tf.newaxis,...]
    y_true_one_hot=tf.one_hot(tf.cast(y_true,tf.uint8),3)
    return tf.reduce_mean(tf.keras.losses.categorical_crossentropy(y_true_one_hot*weights,
                                                                   y_pred,
                                                                   from_logits=False))

I have a function to create initial models in identical states and the model is very simple

def make_model():
    tf.random.set_seed(42)
    np.random.seed(42)
    model=tf.keras.Sequential([
        tf.keras.layers.Dense(3,'softmax',input_shape=[1024,])
    ])
    return model

I instantiate a RMSProp optimizer (code not shown) and call it optimizer and then train the two models so

model1=make_model()
model2=make_model()
model1.compile(loss=customloss1,optimizer=optimizer)
model2.compile(loss=customloss2,optimizer=optimizer)

history1 = model1.fit(x,y,epochs=100,batch_size=50, verbose=0)
history2 = model2.fit(x,y,epochs=100,batch_size=50, verbose=0)

and the results are plotted below

as can be seen the model using the customloss with the decorator doesn't converge after a few epochs. I do not understand what is happening. To dig in to this deeper I decided to do the training loop manually so I made two training step functions that use two copies of the model and differ only in which custom loss function they use

model1=make_model()
model2=make_model()

@tf.function
def train_step1(x,y):

    with tf.GradientTape() as tape:
        predictions  = model1(x)
        loss = customloss1(y, predictions)

    gradients = tape.gradient(loss, model1.trainable_variables)    
    optimizer.apply_gradients(zip(gradients, model1.trainable_variables))
    return loss

@tf.function
def train_step2(x,y):

    with tf.GradientTape() as tape:
        predictions  = model2(x)
        loss = customloss2(y, predictions)

    gradients = tape.gradient(loss, model2.trainable_variables)    
    optimizer.apply_gradients(zip(gradients, model2.trainable_variables))
    return loss

This time the behavior is reversed!

In this case I can understand which one is the right answer by removing all tf.function decorators and running pure python with eager execution and it is the manual training loop with the decorated custom loss. Its results match the keras model fit with class_weight set to the correct values (The curves exactly match).

But I have no idea why the aforementioned combination is giving the correct results. I was under the impression that if the training loop function has a decorator then all the function in the function stack automatically get converted to graph mode and any lower level tf.function decorators are redundant.

Extra Info

I investigated a bit more and turns out that for the custom look the gradients that come out of the two loss functions are different!!!

model3=make_model()

@tf.function
def get_gradients(x,y):
    with tf.GradientTape() as tape1:
        p1=model3(x)
        l1=customloss1(y,p1)
    with tf.GradientTape() as tape2:
        p2=model3(x)
        l2=customloss2(y,p2)

    gradients1=tape1.gradient(l1,model3.trainable_variables)
    gradients2=tape2.gradient(l2,model3.trainable_variables)

    return gradients1, gradients2

gives

([,
  ],
 [,
  ])

The lower one is the correct one and if I remove the decorator on get_gradients then I get the same gradients for both the loss functions. So its some kind of weird interplay of tf.function in the gradient tape it seems.

Decorating a custom loss with tf.function changes the training results completely, both in keras model.fit method as well as custom training loop

0 Answers