Украшение пользовательской потери с помощью tf.function полностью меняет результаты обучения, как в методе keras model.fit, так и в пользовательском цикле обучения

Question

Я сам пытался реализовать class_weight, поскольку я видел какое-то поведение обычным способом (class_weight в методе model.fit), и это привело меня к очень запутанному поведению, и после того, как я подумал об этом довольно долго какое-то время я не могу понять, что происходит.

Весь пост длинный, поэтому я даю краткое описание, прежде чем сообщать подробности тем, кто интересуется

Краткое описание

Вкратце, я определяю две пользовательские функции потерь, которые отличаются только декоратором tf.function. У первого нет декоратора, а у второго есть. При использовании keras model.fit обучение не сходится с тем, в котором есть декоратор. Когда я пробую то же самое с индивидуальным обучением l oop, поведение меняется на противоположное.

Я, очевидно, могу увидеть, что является правильным поведением, запустив все в режиме ожидания, и поэтому я вижу, что пользовательский l oop с функцией украшения является правильным. Его результаты также соответствуют обычному методу подбора модели 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

I создайте экземпляр оптимизатора RMSProp (код не показан) и вызовите его optimizer , а затем обучите две модели, чтобы

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 с декоратором, не сходится через несколько эпох. Я не понимаю, что происходит. Чтобы углубиться в это, я решил провести обучение l oop вручную, поэтому я сделал две ступенчатые функции обучения, которые используют две копии модели и отличаются только тем, какую пользовательскую функцию потерь они используют

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 с нетерпеливым выполнением и это ручной учебный l oop с декорированной индивидуальной потерей. Его результаты соответствуют модели keras, соответствующей правильным значениям class_weight (кривые точно совпадают).

Но я понятия не имею, почему вышеупомянутая комбинация дает правильные результаты. У меня создалось впечатление, что если у обучающей функции l oop есть декоратор, то все функции в стеке функций автоматически преобразуются в режим графика, а декораторы 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 в градиентной ленте.