I'm writing this because as an exercise I'm building my own implementation of YOLO v2 (using Keras with Tensorflow backend). I've read the papers (YOLO and YOLOv2) and I'm having some trouble to completely understand the loss function.
From what I understood, we basically split the loss function into several ones (using MSE). For all the losses we compare only the predicted bbox (made by the anchor with the highest IOU with ground truth) and the real one (or ground truth). But for the confidence loss, we also account for errors made on miss predicted bboxes, (i.e. predicted boxes on background). I leave behind the snippet of my code used for calculating the loss:
def yolo_loss_v2(lambda_coord, lambda_noobj):
# Hack to allow for additional parameters when defining a keras loss function
def calculate_loss(y_true, y_pred):
"""
Computes YOLO Loss
Arguments:
---------
y_pred : [tf.Tensor] Tensor of shape (batch_size, grid_h, grid_w, nb_anchors, 5+nb_classes)
containing predicted values
y_pred : [tf.Tensor] Tensor of shape (batch_size, grid_h, grid_w, nb_anchors, 5+nb_classes)
containing ground truth values
"""
# Convert network's output (see adjust_yolo_ouput for further info)
y_pred = adjust_yolo_output(y_pred)
# Get common values
batch_size = y_true.shape[0]
epsilon = 1e-6 # Epsilon to avoid division by 0
# Get mask for indexing where there're objects detected on ground truth
# By the way we've created the batch, we already calculated which prior (anchor box)
# has the highest IOU with the real object
responsability_selector = tf.where(y_true[..., 4] == 1, 1, 0)
nb_bboxes = tf.reduce_sum(responsability_selector, [1,2,3])
nb_bboxes = tf.cast(nb_bboxes, tf.float32)
responsability_selector = tf.cast(responsability_selector, tf.float32)
# ------------------------------------------------------------------------
# Localizaton LOSS (i.e. differences in true bboxes and the predicted ones)
# Calculate x,y loss
pred_xy = y_pred[..., 0:2] * responsability_selector[..., None]
true_xy = y_true[..., 0:2] * responsability_selector[..., None]
loss_xy = tf.reduce_sum(tf.math.square(pred_xy - true_xy), [1,2,3,4]) / (nb_bboxes + epsilon)
loss_xy = tf.reduce_sum(loss_xy, 0) / batch_size
# Calculate w,h loss
pred_wh = y_pred[..., 2:4] * responsability_selector[..., None]
true_wh = y_true[..., 2:4] * responsability_selector[..., None]
# print('{}\nDebug:\nloss_wh (min, max):{:.2f} {:.2f}'.format("*"*15,
# tf.reduce_min(pred_wh), tf.reduce_max(pred_wh)))
pred_wh = tf.math.sqrt(pred_wh)
true_wh = tf.math.sqrt(true_wh)
loss_wh = tf.reduce_sum(tf.math.square(pred_wh - true_wh), [1,2,3,4]) / (nb_bboxes + epsilon)
loss_wh = tf.reduce_sum(loss_wh, 0) / batch_size
# ------------------------------------------------------------------------
# Confidence LOSS (i.e. differences between anchor confidence vs IOU scores
# of predictions and ground truth )
pred_c_obj = y_pred[..., 4] * responsability_selector
true_c_obj = y_true[..., 4] * responsability_selector
loss_c = tf.reduce_sum(tf.math.square(pred_c_obj - true_c_obj), [1,2,3]) / (nb_bboxes + epsilon)
loss_c = tf.reduce_sum(loss_c, 0) / batch_size
pred_c_noobj = y_pred[..., 4] * (responsability_selector - 1)
true_c_noobj = y_true[..., 4] * (responsability_selector - 1)
loss_c_noobj = tf.reduce_sum(tf.math.square(pred_c_noobj - true_c_noobj), [1,2,3]) / (nb_bboxes + epsilon)
loss_c_noobj = lambda_noobj * tf.reduce_sum(loss_c_noobj, 0) / batch_size
loss_c += loss_c_noobj
# ------------------------------------------------------------------------
# Classification LOSS (i.e. class probabilities discrepancies)
pred_class = y_pred[..., 5:] * responsability_selector[..., None]
true_class = y_true[..., 5:] * responsability_selector[..., None]
loss_class = tf.reduce_sum(tf.math.square(pred_class - true_class), [1,2,3,4]) / (nb_bboxes + epsilon)
loss_class = tf.reduce_sum(loss_class, 0) / batch_size
# ------------------------------------------------------------------------
# Total LOSS
# print('{}\nDebug:\nloss_xy: {}\nloss_wh:{}\nloss_c:{}\nloss_class:{}'.format("*"*15,
# loss_xy, loss_wh, loss_c, loss_class))
return lambda_coord*(loss_xy + loss_wh) + loss_c + loss_class
return calculate_loss
My y_true was created using a batch generator and has a shape equal to (BATCH_SIZE, NB_GRIDS_Y, NB_GRIDS_X, NB_ANCHORS, 5+NB_CLASSES) and adjust_yolo_output basically transforms the network's output into the real quantities we're predicting. As stated in the paper, the network's output for each anchor (for each gridcell) is a bbox (tx, ty, tw, th, to) such that (quoting YOLOv2 paper):
bx = sigmoid(tx) + cx
by = sigmoid(ty) + cy
bw = pw*exp(tw)
bh = ph*exp(th)
Pr(obj)*IOU(b,obj) = sigmoid(to)
Pr(Ci) = softmax([C1 C2 .. Cn])
where (cx, cy) is the gridcell's origin and (ph, pw) are the corresponding anchor dimensions.
The hardest part for me to understand is Pr(obj)*IOU(b,obj) = sigmoid(to). When I created my batch, I've already calculated which anchor has the highest IOU with the labeled bbox which corresponds to my axis=3 (the index in this axis indicates which anchor I'm using). Shouldn't I want sigmoid(to) to be equal to the probability that there's an object present in that gridcell (with a similar shape as the indexed anchor) Pr(obj) alone, so that when I want to calculate the probability for each class, I'd just multiply this by the conditional probability P(class i | obj) predicted by the model. But instead I'll be getting the probability of an object present multiplied by the IOU score? I mean, I could divide sigmoid(to) by the IOU score on training data, but how would that work when making a prediction on unseen (unlabeled) data?
Quoting the original paper (YOLO):
We assign one predictor to be “responsible” for predicting an object based on which prediction has the highest current IOU with the ground truth. This leads to specialization between the bounding box predictors.
I should remark that I'm implementing YOLO v2, but as there's no mention of modifications to the loss function, I assumed it uses the same function as the original YOLO architecture.
Any other mistake you find (or feedback you have) will be greatly appreciated. Please forgive me if I've made some basic mistakes (it's both my first "real" code using TensorFlow and question asked here). Tell me if I did not make myself clear.
Thank you all for your time!