keras model.fit() 和 TF tape.gradient() 给出不同的结果

Question

I have a model that I am building using the keras functional API. After defining it, I compile it with the SGD optimizer as follows.

opt = tf.keras.optimizers.SGD(learning_rate=0.05, momentum=0.9, decay=1e-3,clipnorm=1)
model.compile(optimizer=opt, loss='mse')
model.fit(train_datagen, epochs=50,shuffle=True,verbose=True)

This works fine and my model converges as expected.

However, when I to implement the same exact functionality using TF's tape gradient, I consistently get NaN gradients which cause my weights to equate to NaN and subsequently my loss function value because NaN. Here is the code that I use:

opt = tf.keras.optimizers.SGD(learning_rate=0.05, momentum=0.9, decay=1e-3,clipnorm=1)
loss_fn = tf.keras.losses.MeanSquaredError()

epochs = 50

for epoch in range(epochs):
        
    batch_list = list(range(len(train_datagen)))
    random.shuffle(batch_list)

    running_loss = 0
    
    for ii in batch_list:
        x,y_true = train_datagen[ii]
        
        with tf.GradientTape() as tape:
            y_pred = model(x, training=True)
            loss_value = loss_fn(y_true,y_pred)
            
        grads = tape.gradient(loss_value, model.trainable_variables)
        opt.apply_gradients(zip(grads, model.trainable_variables))
        
        running_loss += loss_value
    
    print('Epoch',epoch,'Running Loss:',running_loss.numpy()/len(batch_list))

Is the code that I wrote equivalent to the Keras model.fit() functionality? For some reason, when I use the above code, I consistently get NaN gradients but with model.fit() it never happens.

Answer 1

I think this is because the class based loss, MeanSquaredError , requires some extra tinkering to get it working outside of the .fit method. Instead, use the functional one to make it easier. Just call it inside the training step like this.

with tf.GradientTape() as tape:
    y_pred = model(x, training=True)
    loss_value = tf.keras.losses.mean_squared_error(y_true,y_pred)