Keras-yolo3: Loss components as train metric to be displayed in Tensorboard

[AlbertoSabater]

Hi! Let's say that we can decompose the training loss in four components: xy_loss, wh_loss, confidence_loss and class_loss.
Is there any way to calculate these metrics to be plotted in Tensorboard?
The motivation of this idea is that this loss implementation doesn't include parameters to weight each one of these loss components. So, by visualizing their values we can check if those weights should be worthy to implement. And they would also be useful to compare how proficient are different models in each of the Yolo tasks.

The problem is that a Keras metric function only can accept y_true and y_pred as parameters but the metrics I propose also needs to take the anchors and num_classes as parameters.

[AlbertoSabater]

Finally, I've managed to add the loss components as metrics during training. To do that, I had to modify the yolo_loss function to return the loss along with each one of its components:

def yolo_loss(args, anchors, num_classes, ignore_thresh=.5):
    '''Return yolo_loss tensor composed by the loss and all its components

    Parameters
    ----------
    yolo_outputs: list of tensor, the output of yolo_body or tiny_yolo_body
    y_true: list of array, the output of preprocess_true_boxes
    anchors: array, shape=(N, 2), wh
    num_classes: integer
    ignore_thresh: float, the iou threshold whether to ignore object confidence loss

    Returns
    -------
    loss: tensor, shape=(1,)

    '''
    num_layers = len(anchors)//3 # default setting
    yolo_outputs = args[:num_layers]
    y_true = args[num_layers:]
    anchor_mask = [[6,7,8], [3,4,5], [0,1,2]] if num_layers==3 else [[3,4,5], [1,2,3]]
    input_shape = K.cast(K.shape(yolo_outputs[0])[1:3] * 32, K.dtype(y_true[0]))
    grid_shapes = [K.cast(K.shape(yolo_outputs[l])[1:3], K.dtype(y_true[0])) for l in range(num_layers)]
    loss = 0
    total_xy_loss, total_wh_loss, total_confidence_loss_obj, total_confidence_loss_noobj, total_class_loss = 0, 0, 0, 0, 0
    m = K.shape(yolo_outputs[0])[0] # batch size, tensor
    mf = K.cast(m, K.dtype(yolo_outputs[0]))

    for l in range(num_layers):
        object_mask = y_true[l][..., 4:5]
        true_class_probs = y_true[l][..., 5:]

        grid, raw_pred, pred_xy, pred_wh = yolo_head(yolo_outputs[l],
             anchors[anchor_mask[l]], num_classes, input_shape, calc_loss=True)
        pred_box = K.concatenate([pred_xy, pred_wh])

        # Darknet raw box to calculate loss.
        raw_true_xy = y_true[l][..., :2]*grid_shapes[l][::-1] - grid
        raw_true_wh = K.log(y_true[l][..., 2:4] / anchors[anchor_mask[l]] * input_shape[::-1])
        raw_true_wh = K.switch(object_mask, raw_true_wh, K.zeros_like(raw_true_wh)) # avoid log(0)=-inf
        box_loss_scale = 2 - y_true[l][...,2:3]*y_true[l][...,3:4]

        # Find ignore mask, iterate over each of batch.
        ignore_mask = tf.TensorArray(K.dtype(y_true[0]), size=1, dynamic_size=True)
        object_mask_bool = K.cast(object_mask, 'bool')
        def loop_body(b, ignore_mask):
            true_box = tf.boolean_mask(y_true[l][b,...,0:4], object_mask_bool[b,...,0])
            iou = box_iou(pred_box[b], true_box)
            best_iou = K.max(iou, axis=-1)
            ignore_mask = ignore_mask.write(b, K.cast(best_iou<ignore_thresh, K.dtype(true_box)))
            return b+1, ignore_mask
        _, ignore_mask = K.control_flow_ops.while_loop(lambda b,*args: b<m, loop_body, [0, ignore_mask])
        ignore_mask = ignore_mask.stack()
        ignore_mask = K.expand_dims(ignore_mask, -1)

        # K.binary_crossentropy is helpful to avoid exp overflow.
        xy_loss = object_mask * box_loss_scale * K.binary_crossentropy(raw_true_xy, raw_pred[...,0:2], from_logits=True)
        wh_loss = object_mask * box_loss_scale * 0.5 * K.square(raw_true_wh-raw_pred[...,2:4])
        confidence_loss_obj = object_mask * K.binary_crossentropy(object_mask, raw_pred[...,4:5], from_logits=True)
        confidence_loss_noobj = (1-object_mask) * K.binary_crossentropy(object_mask, raw_pred[...,4:5], from_logits=True) * ignore_mask
#       confidence_loss = object_mask * K.binary_crossentropy(object_mask, raw_pred[...,4:5], from_logits=True)+ \
#           (1-object_mask) * K.binary_crossentropy(object_mask, raw_pred[...,4:5], from_logits=True) * ignore_mask
        class_loss = object_mask * K.binary_crossentropy(true_class_probs, raw_pred[...,5:], from_logits=True)

        xy_loss = K.sum(xy_loss) / mf
        wh_loss = K.sum(wh_loss) / mf
        confidence_loss_obj = K.sum(confidence_loss_obj) / mf
        confidence_loss_noobj = K.sum(confidence_loss_noobj) / mf
#       confidence_loss = K.sum(confidence_loss) / mf
        class_loss = K.sum(class_loss) / mf

        total_xy_loss += xy_loss
        total_wh_loss += wh_loss
#       total_confidence_loss += confidence_loss
        total_confidence_loss_obj += confidence_loss_obj
        total_confidence_loss_noobj += confidence_loss_noobj
        total_class_loss += class_loss

    loss += total_xy_loss + total_wh_loss + total_confidence_loss_obj + total_confidence_loss_noobj + total_class_loss
#   if print_loss:
#       loss = tf.Print(loss, [loss, xy_loss, wh_loss, confidence_loss, class_loss, K.sum(ignore_mask)], message='loss: ')
#   return loss
    return tf.convert_to_tensor([loss, 
                              total_xy_loss, total_wh_loss, 
                              total_confidence_loss_obj + total_confidence_loss_noobj,
                              total_confidence_loss_obj, total_confidence_loss_noobj, 
                              total_class_loss], dtype=tf.float32)

And define the functions to extract each value:

def loss(y_true, y_pred): return y_pred[0]
def xy_loss(y_true, y_pred): return y_pred[1]
def wh_loss(y_true, y_pred): return y_pred[2]
def confidence_loss(y_true, y_pred): return y_pred[3]
def confidence_loss_obj(y_true, y_pred): return y_pred[4]
def confidence_loss_noobj(y_true, y_pred): return y_pred[5]
def class_loss(y_true, y_pred): return y_pred[6]

Finally, you have to add these functions to the model compilation:

model.compile(optimizer = optimizer,
              loss = {'yolo_loss': loss},       # use custom yolo_loss Lambda layer.
              metrics = [
                         xy_loss, wh_loss, confidence_loss, 
                         confidence_loss_obj, confidence_loss_noobj, class_loss,
                         train_utils.get_lr_metric(optimizer),
                         ]  
              )

[zhounan95]

请问train_utils.get_lr_metric(optimizer)在哪儿调用呀?