使用 tf.function 装饰自定义损失会完全改变训练结果，无论是在 keras model.fit 方法还是自定义训练循环中

Question

我试图自己实现 class_weight 因为我以通常的方式（model.fit 方法中的 class_weight）看到了一些行为，这让我看到了一个非常令人困惑的行为，经过一段时间的思考后我无法理解正在发生的事情。

整篇文章很长，所以我先简要说明一下，然后再为感兴趣的人提供详细信息

简要描述;简介

简而言之，我定义了两个自定义损失函数，它们仅在 tf.function 装饰器中有所不同。第一个没有装饰器，而第二个有。使用 keras model.fit，训练不会与具有装饰器的训练收敛。当我对自定义训练循环进行相同尝试时，行为会逆转。

通过在 Eager 模式下运行所有​​内容，我可以明显看出哪种行为是正确的，因此我可以看到带有装饰函数的自定义循环是正确的。它的结果也与通常的 keras 模型拟合方法与 class_weight 匹配。所有其他三种方法都会给出错误的结果，但我不明白为什么。

详细说明

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

我有一个功能可以在相同的状态下创建初始模型，并且模型非常简单

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

我实例化一个 RMSProp 优化器（代码未显示）并将其称为优化器，然后训练这两个模型

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)

结果如下图所示

可以看出，使用带有装饰器的 customloss 的模型在几个时期后没有收敛。我不明白发生了什么。为了深入研究这个问题，我决定手动进行训练循环，所以我制作了两个训练步骤函数，它们使用模型的两个副本，不同之处仅在于它们使用的自定义损失函数

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

这一次的行为是相反的！

在这种情况下，我可以通过删除所有 tf​​.function 装饰器并通过急切执行运行纯 python 来理解哪一个是正确的答案，这是带有装饰自定义损失的手动训练循环。它的结果与 keras 模型匹配，class_weight 设置为正确的值（曲线完全匹配）。

但我不知道为什么上述组合给出了正确的结果。我的印象是，如果训练循环函数有一个装饰器，那么函数堆栈中的所有函数都会自动转换为图形模式，并且任何较低级别的 tf.function 装饰器都是多余的。

额外信息

我进行了更多调查，结果发现对于自定义外观，两个损失函数产生的梯度是不同的！！！

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

给

([<tf.Tensor: shape=(1024, 3), dtype=float32, numpy=
  array([[-0.01336379,  0.10262163, -0.01915502],
         [ 0.07654451, -0.0181675 , -0.04819181],
         [ 0.00367431, -0.06802277,  0.05872081],
         ...,
         [-0.13633026,  0.01184574,  0.02273583],
         [ 0.02155258, -0.04340569,  0.0841853 ],
         [ 0.17315787, -0.12444994, -0.04137734]], dtype=float32)>,
  <tf.Tensor: shape=(3,), dtype=float32, numpy=array([0.05435816, 0.02056887, 0.28507292], dtype=float32)>],
 [<tf.Tensor: shape=(1024, 3), dtype=float32, numpy=
  array([[-0.00059669,  0.06096834, -0.06037165],
         [ 0.03078352,  0.00224353, -0.03302703],
         [ 0.01101062, -0.04670432,  0.03569368],
         ...,
         [-0.12952083,  0.06808907,  0.06143178],
         [ 0.03497609, -0.09745409,  0.06247801],
         [ 0.131704  , -0.12800933, -0.00369466]], dtype=float32)>,
  <tf.Tensor: shape=(3,), dtype=float32, numpy=array([-0.05939803, -0.10304251,  0.16244057], dtype=float32)>])

较低的一个是正确的，如果我删除 get_gradients 上的装饰器，那么我会得到两个损失函数的相同梯度。所以它似乎在梯度磁带中出现了某种奇怪的 tf.function 相互作用。