The beginner’s guide to implementing YOLOv3 in TensorFlow 2.0 (part-3)

In part 2, we’ve discovered how to construct the YOLOv3 network. In this part 3, we’ll focus on the file yolov3.weights.

So, what we’re going to do in part is to load the weights parameters from the file yolov3.weights, then convert them into the TensorFlow 2.0 weights format. Just to remain you that, the file yolov3.weights contains the pre-trained CNN’s parameters of YOLOv3.

To begin with, let’s take a look at how the YOLOv3 weights are stored.

How YOLOv3’s weights are stored?

The original YOLOv3 weights file yolov3.weights is a binary file and the weights are stored in the float data type.

One thing that we need to know that the weights only belong to convolutional layers. As we know, in YOLOv3, there are 2 convolutional layer types, with and without a batch normalization layer. So, the weights are applied differently for different types of convolutional layers. Since we’re reading the only float data, there’s no clue which one belongs to which layer. If we incorrectly associate these weights with their layers properly, we’ll screw up everything, and the weights won’t be converted. So, that’s why understanding how the weights are stored is crucially important.

Here, I tried to make a simple flowchart in order to describe how the weights are stored.

When we re-write these weights to TensorFlow’s format for a convolutional with a batch normalization layer, we need to switch the position of beta and gamma. So, they’re ordered like this: beta, gamma, means, variance and conv weights. However, the weights’ order remains the same for the convolutional without a batch normalization layer.

All right!!, Now we’re ready to code the weights converter.

Without further ado, let’s do it…

Working with the file convert_weights.py

Open the file convert_weights.py, then copy and paste the following code to the top of it. Here, we import NumPy library and the two functions that we’ve created previously in part 2, YOLOv3Net and parse_cfg.

#convert_weights.py
import numpy as np
from yolov3 import YOLOv3Net
from yolov3 import parse_cfg

Now, let’s create a function called load_weights(). This function has 3 parameters, model, cfgfile, and weightfile. The parameter model is a returning parameters of the network’s model after calling the function YOLOv3Net. Thecfgfile and weightfile are respectively refer to the files yolov3.cfg and yolov3.weights.

def load_weights(model,cfgfile,weightfile):

Open the file yolov3.weights and read the first 5 values. These values are the header information. So, we can skip them all.

    # Open the weights file
    fp = open(weightfile, "rb")

    # The first 5 values are header information
    np.fromfile(fp, dtype=np.int32, count=5)

Then call parse_cfg() function.

    blocks = parse_cfg(cfgfile)

As we did when building the YOLOv3 network, we need to loop over the blocks and search for the convolutional layer. Don’t forget to check whether the convolutional is with batch normalization or not. If it is true, go get the relevant values (gamma, beta, means, and variance), and re-arrange them to the TensorFlow weights order. Otherwise, take the bias values. After that take the convolutional weights and set these weights to the convolutional layer depending on the convolutional type.

    for i, block in enumerate(blocks[1:]):

        if (block["type"] == "convolutional"):
            conv_layer = model.get_layer('conv_' + str(i))
            print("layer: ",i+1,conv_layer)

            filters = conv_layer.filters
            k_size = conv_layer.kernel_size[0]
            in_dim = conv_layer.input_shape[-1]

            if "batch_normalize" in block:

                norm_layer = model.get_layer('bnorm_' + str(i))
                print("layer: ",i+1,norm_layer)
                size = np.prod(norm_layer.get_weights()[0].shape)

                bn_weights = np.fromfile(fp, dtype=np.float32, count=4 * filters)
                # tf [gamma, beta, mean, variance]
                bn_weights = bn_weights.reshape((4, filters))[[1, 0, 2, 3]]

            else:
                conv_bias = np.fromfile(fp, dtype=np.float32, count=filters)

            # darknet shape (out_dim, in_dim, height, width)
            conv_shape = (filters, in_dim, k_size, k_size)
            conv_weights = np.fromfile(
                fp, dtype=np.float32, count=np.product(conv_shape))
            # tf shape (height, width, in_dim, out_dim)
            conv_weights = conv_weights.reshape(
                conv_shape).transpose([2, 3, 1, 0])

            if "batch_normalize" in block:
                norm_layer.set_weights(bn_weights)
                conv_layer.set_weights([conv_weights])
            else:
                conv_layer.set_weights([conv_weights, conv_bias])

Alert if the reading has failed. Then, close the file whether the reading was successful or not.

    assert len(fp.read()) == 0, 'failed to read all data'
    fp.close()

The last part of this code is the main function. Copy and paste the following code of the main function just right after the function load_weights().

def main():

    weightfile = "weights/yolov3.weights"
    cfgfile = "cfg/yolov3.cfg"

    model_size = (416, 416, 3)
    num_classes = 80

    model=YOLOv3Net(cfgfile,model_size,num_classes)
    load_weights(model,cfgfile,weightfile)

    try:
        model.save_weights('weights/yolov3_weights.tf')
        print('\nThe file \'yolov3_weights.tf\' has been saved successfully.')
    except IOError:
        print("Couldn't write the file \'yolov3_weights.tf\'.")

Here is the complete code for the convert_weights.py

#convert_weights.py
import numpy as np
from yolov3 import YOLOv3Net
from yolov3 import parse_cfg

def load_weights(model,cfgfile,weightfile):
    # Open the weights file
    fp = open(weightfile, "rb")

    # Skip 5 header values
    np.fromfile(fp, dtype=np.int32, count=5)

    # The rest of the values are the weights
    blocks = parse_cfg(cfgfile)

    for i, block in enumerate(blocks[1:]):

        if (block["type"] == "convolutional"):
            conv_layer = model.get_layer('conv_' + str(i))
            print("layer: ",i+1,conv_layer)

            filters = conv_layer.filters
            k_size = conv_layer.kernel_size[0]
            in_dim = conv_layer.input_shape[-1]

            if "batch_normalize" in block:

                norm_layer = model.get_layer('bnorm_' + str(i))
                print("layer: ",i+1,norm_layer)
                size = np.prod(norm_layer.get_weights()[0].shape)

                bn_weights = np.fromfile(fp, dtype=np.float32, count=4 * filters)
                # tf [gamma, beta, mean, variance]
                bn_weights = bn_weights.reshape((4, filters))[[1, 0, 2, 3]]

            else:
                conv_bias = np.fromfile(fp, dtype=np.float32, count=filters)

            # darknet shape (out_dim, in_dim, height, width)
            conv_shape = (filters, in_dim, k_size, k_size)
            conv_weights = np.fromfile(
                fp, dtype=np.float32, count=np.product(conv_shape))
            # tf shape (height, width, in_dim, out_dim)
            conv_weights = conv_weights.reshape(
                conv_shape).transpose([2, 3, 1, 0])

            if "batch_normalize" in block:
                norm_layer.set_weights(bn_weights)
                conv_layer.set_weights([conv_weights])
            else:
                conv_layer.set_weights([conv_weights, conv_bias])

    assert len(fp.read()) == 0, 'failed to read all data'
    fp.close()


def main():

    weightfile = "weights/yolov3.weights"
    cfgfile = "cfg/yolov3.cfg"

    model_size = (416, 416, 3)
    num_classes = 80

    model=YOLOv3Net(cfgfile,model_size,num_classes)
    load_weights(model,cfgfile,weightfile)

    try:
        model.save_weights('weights/yolov3_weights.tf')
        print('\nThe file \'yolov3_weights.tf\' has been saved successfully.')
    except IOError:
        print("Couldn't write the file \'yolov3_weights.tf\'.")


if __name__ == '__main__':
    main()

Finally, we can now execute the convert_weights.py file . Open Anaconda prompt or terminal in Pycharm, type the following command and press Enter.

python convert_weights.py

Here is the output, I printed all the convolutional layers just to make sure that the weights are loaded correctly until the last convolutional layer.

If you use PyCharm, look at the Project Navigation on the left side as I pointed with the red boxes, you have 4 new files:

• checkpoint
• yolov3_weights.tf.data-00000-of-00002
• yolov3_weights.tf.data-00001-of-00002
• yolov3_weights.tf.index

Those files are the TensorFlow 2.0 weights format. So, anytime we want to use them, just simply call them like the only one file, yolov3_weights.tf. We’ll see how to do this in the last part of this tutorial.

This is the end of part 3, and we still have several things to do, those are:

• creating a pipeline to read the input image or video/camera,
• computing the prediction,
• and drawing the bounding boxes prediction over the input image/video/camera.

Those are what we will be doing soon in the last part. Let’s go into it…

Parts

• Part-1, An introduction of the YOLOv3 algorithm.

• Part-2, Parsing the YOLOv3 configuration file and creating the YOLOv3 network.

• Part-3, Converting the YOLOv3 pre-trained weights into the TensorFlow 2.0 weights format.

• Part-4, Encoding bounding boxes and testing this implementation with images and videos.