Пользовательские потери Keras дают разные результаты при тренировках

Question

Я создал пользовательскую функцию потерь для модели YOLO в Керасе. Тем не менее, я получаю другие результаты с функцией потери во время тренировки, чем когда я использую ее вне тренировки.

Вот моя пользовательская потеря:

@tf.function 
def yolo_loss(y_pred,y_true):

    pred_boxes = y_pred[...,:4]
    true_boxes = y_true[...,:4]

    # use boxes to get proper xy and wh as from the paper
    pred_cx,pred_cy = get_all_centers(pred_boxes)
    true_cx,true_cy = get_all_centers(true_boxes)

    pred_w,pred_h = get_wh(pred_boxes)
    true_w,true_h = get_wh(true_boxes)

    # create masks
    pred_confs = y_pred[...,4]
    true_confs = y_true[...,4] # this is equivalent to 1obj from the paper... it will be 1 if there is supposed to be something there


    obj = true_confs.__copy__()
    noobj = tf.where(obj==0,1,0)
    noobj = tf.cast(noobj,tf.float32)

    pred_clss  = y_pred[...,5:]
    true_clss  = y_true[...,5:]

    xy_loss = K.sum(obj*(
            ((pred_cx-true_cx)**2)
            +
            ((pred_cy-true_cy)**2)))

    wh_loss = K.sum(obj*(((pred_w-true_w)**2)+
                         ((pred_h-true_h)**2)))

    loc_loss = xy_loss+wh_loss

    conf_loss =K.sum(obj*(pred_confs-true_confs)**2)\
               +\
               K.sum(noobj*(pred_confs-true_confs)**2)


    cls_loss = K.sum(obj*(K.sum((pred_clss-true_clss)**2,axis=-1)))

    loss = loc_loss+conf_loss+cls_loss


    tf.print(xy_loss)
    tf.print(wh_loss)
    tf.print(conf_loss)
    tf.print(cls_loss)

    return loss

4 числа представляют различные части функции потери. Код доступен здесь: https://github.com/GoldwinXS/custom_yolo

выходные данные выглядят так (я удалил часть текста до того, как возникла проблема):

640/1000 [= =================> ...........] - ETA: 8 с - потеря: 237.2000

0

0

231

0

768/1000 [=====================>. ......] - ETA: 5 с - потеря: 236,1667

0

0

235

0

896 / 1000 [=========================> ....] - ETA: 2 с - потеря: 236,0000

0

0

194

0

1000/1000 [========================== ==========] - 25 с 25 мс / шаг - потеря: 231,6320

, а затем при вызове моей потери после тренировки я получаю:

1561,9729

694.685425

1846

9304.88672

Это можно воспроизвести, запустив MyYoloTest.py

Кто-нибудь знает, почему это может происходить? ?

РЕДАКТИРОВАТЬ: После большой боли я обнаружил, что y_pred входит в функцию потерь с формой TensorShape ([None, None, None, None]), тогда как y_true имеет ожидаемую форму: TensorShape ([None, 3 , 3, 8])

У кого-нибудь есть идеи, почему это так?