architecture.py 文件源码

def yolo_loss(labels, predictions, mask):
    masked_labels = tf.boolean_mask(labels, mask)
    masked_predictions = tf.boolean_mask(predictions, mask)

    # ious = tensor_iou(masked_predictions[..., 1:5], masked_labels[..., 1:5])
    # ious = tf.expand_dims(ious, axis=-1)

    xy_loss = tf.reduce_sum((masked_labels[..., :2] - masked_predictions[..., 1:3]) ** 2)
    wh_loss = tf.reduce_sum((tf.sqrt(masked_predictions[..., 3:5]) - tf.sqrt(masked_labels[..., 2:4])) ** 2)

    #     conf_loss = tf.reduce_sum((masked_predictions[..., 0] - ious) ** 2)

    conf_loss = tf.reduce_sum((1 - masked_predictions[..., 0]) ** 2)

    no_obj_loss = tf.reduce_sum((tf.boolean_mask(predictions, ~mask)[..., 0] ** 2))

    class_loss = tf.reduce_sum((masked_predictions[..., 5:] - masked_labels[..., 4:]) ** 2)

    loss = 5 * (xy_loss + wh_loss) + conf_loss + no_obj_loss + class_loss

    return loss