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

Question

我有一个使用keras函数API构建的模型。在定义它之后，我使用SGD优化器编译它，如下所示。

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)

这可以很好地工作，并且我的模型会像预期的那样收敛。

然而，当我使用TF的磁带梯度实现相同的功能时，我总是得到NaN梯度，这导致我的权重等于NaN，随后我的损失函数值因为NaN而相等。下面是我使用的代码：

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))

我写的代码是否等同于Keras的model.fit()功能？由于某种原因，当我使用上面的代码时，我总是得到NaN渐变，但使用model.fit()时，这种情况永远不会发生。

Answer 1

我认为这是因为基于类的损失MeanSquaredError需要在.fit方法之外进行一些额外的修补工作。取而代之的是，使用函数式的来使它更容易。就像这样在训练步骤中调用它。

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