How to train the predicting boxes in a YOLO network?

Question

I have just finished this tutorial that explains how YOLO networks work. Instead of training the network's weights with a training set, the author loads pre-trained weights and uses them to test the network. I'm interested in how you would train this network.

I want to run the network on an image that is size 1000 x 1000 and can hold objects with 10 different classes. The network's forward function returns a tensor that has dimensions

[1, num_bounding_boxes * final_layer_grid_width * final_layer_grid_height, 15].

For each cell and bounding box in the final layer, the network predicts 15 values:

x_center, y_center, bbox_width, bbox_height, object_confidence_score, class_1_confidence_score, ..., class_10_confidence_score

I have a list of labels corresponding to the images. Each element in the list (corresponding to an image) looks like this:

[[313,567,47, 23, 4],
[398, 122, 57, 32, 6],
 ...,
[499, 993, 47, 19, 8]]

The first element indicates that the point (313, 567) of the original H x W image is the center of an object, that the corresponding bounding box is of width 47, of height 23, and that the class is 4.

My plan was to check to which grid cell in the final layer each object should be assigned based on its coordinates and the stride of last layer. Suppose that the last layer is 100 x 100 and the input image is 1000 x 1000 then the first element in the example label should be predicted by grid cell [31,56], because the stride is 10. I think that I can create a target Tensor based on this method and I think that I can train the object confidence score and the 10 different class scores this way.

What I'm not quite sure of is how to train the x_center, y_center, bbox_width, bbox_height. What values should these parameters hold if the grid cell in the final layer should not predict a bounding box?

Question: How do I train the x_center, y_center, bbox_width, bbox_height parameters?

Answer 1

This is a good question and recaps the thinking of how YOLO works.

So first of all you say your forward path gives you a tensor of

[1, num_bounding_boxes * final_layer_grid_width * final_layer_grid_height, 15].

For simplicity let us reshape this to [grid_width, grid_height, num_boxes, x_center, y_center, bbox_width, bbox_height, object_confidence_score, class1, ..., class10], this means for each Gridcell and each Boundingboxprior you have one value for x_center, y_center, bbox_width, bbox_height.

x_center and y_center predicts a value which is sigmoid activated the offset from gridcell border to the actual midpoint of the bounding box. So b_x = sigmoid(x_center)+ c_x, where c_x is x coordinate of grid Cell, for b_y = sigmoid(y_center)+ c_y, accordingly.

bbox_width and bbox_height are values which are shifts of the BBoxprior in log space so BB_width = BB_prior*e^(bbox_width), for BB_height respectively.

You can think about it this way: If an object is in a grid cell, the object_confidence_score will tell you. So you just have to adjust the Boundingbox. First step is to calculate the bounding box with highest IOU, next step is to adjust it's center with x_center, y_center and the height/widht with bbox_width, bbox_height.

Check the Yolo-V3 Paper for the used formulas and especially figure 2 depicts the parameters you asked for.