Effective conversion of YOLOv3 model to optimized .tflite model

5 min readFeb 28, 2021

Once you are done building an object detection model (using YOLOv3), you take a step forward to use it on edge devices, which would require comparatively optimized or smaller models.

There are several ways of converting YOLO .weights and .cfg files to an optimized .tflite file, such as:

II.

III.

However, this article will explore only one of them (the 1st one) that worked out the best for my object detection project.

Note: Not sure if it would work for other versions of YOLO except YOLOv3.

The procedure is as follows:

Setting up the environment in Google Colab.
Converting .weights & .cfg files to .pb file
Obtaining names of an input layer and three output layers
Converting .pb and other files into a frozen graph (‘.pb file’)
Converting the frozen graph model into an optimized TFlite model

The platform used: Google Colab.

Requirement: YOLOv3 .weights and .cfg files

Step I: Setting up the environment in Google Colab.

Install the 1.15.2 version of Tensorflow before proceeding further.

Connect the Colab Notebook with Google drive.

Step II: Converting .weights and .cfg files to .pb file

Change the working directory and make a new directory ‘New_folder’.

Now, change the directory to ‘New_folder’ and clone this repository - DW2TF.

Go to the folder ‘/New_folder/DW2TF/data’ in your drive and upload .weights and .cfg files in it.

Go to the Colab notebook again and change the directory to DW2TF.

Run this code. It takes .weights and .cfg files as inputs. (Don’t forget to change the names of .weights and .cfg files)

It should output: .pb file, .ckpt.meta, .ckpt.index, .ckpt.data-00000-of-00001, checkpoint in the ‘data’ folder.

Step III: Obtaining the names of an input layer and three output layers

Before jumping onto the next step of converting the .pb file into a frozen graph, we need to obtain the names of an input layer and three output layers that will be used later. There are two ways to do the same.

Either

a. Run this code:

It should print the names of all the layers in the model. The 1st line must have the input layer's name (for instance- ‘yolov3/net1’). For getting YOLOv3 output layers’ names, search for names ending with ‘BiasAdd’ (for instances- yolov3/convolutional59/BiasAdd’, yolov3/convolutional67/BiasAdd, yolov3/convolutional75/BiasAdd)

b. Install ‘Netron’ from this link.

Download the .pb file from the ‘data’ folder on your PC and open it with ‘Netron’ software to visualize the entire model. Either Press ‘Ctrl + Enter’ (for Windows) to get all the names or find the model's terminal points by simply dragging the mouse around.

Step IV: Converting .pb and other files into a frozen graph (‘.pb file’)

Copy this code, save it as ‘to_frozen_graph.py’ and upload it in the ‘DW2TF’ folder in your drive.

import os, argparse

import tensorflow as tf

# The original freeze_graph function
# from tensorflow.python.tools.freeze_graph import freeze_graph 

dir = os.path.dirname(os.path.realpath(__file__))

def freeze_graph(model_dir, output_node_names):
    """Extract the sub graph defined by the output nodes and convert 
    all its variables into constant 
    Args:
        model_dir: the root folder containing the checkpoint state file
        output_node_names: a string, containing all the output node's names, 
                            comma separated
    """
    if not tf.gfile.Exists(model_dir):
        raise AssertionError(
            "Export directory doesn't exists. Please specify an export "
            "directory: %s" % model_dir)

    if not output_node_names:
        print("You need to supply the name of a node to --output_node_names.")
        return -1

    # We retrieve our checkpoint fullpath
    checkpoint = tf.train.get_checkpoint_state(model_dir)
    input_checkpoint = checkpoint.model_checkpoint_path
    
    # We precise the file fullname of our freezed graph
    absolute_model_dir = "/".join(input_checkpoint.split('/')[:-1])
    output_graph = absolute_model_dir + "/frozen_model.pb"

    # We clear devices to allow TensorFlow to control on which device it will load operations
    clear_devices = True

    # We start a session using a temporary fresh Graph
    with tf.Session(graph=tf.Graph()) as sess:
        # We import the meta graph in the current default Graph
        saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=clear_devices)

        # We restore the weights
        saver.restore(sess, input_checkpoint)

        # We use a built-in TF helper to export variables to constants
        output_graph_def = tf.graph_util.convert_variables_to_constants(
            sess, # The session is used to retrieve the weights
            tf.get_default_graph().as_graph_def(), # The graph_def is used to retrieve the nodes 
            output_node_names.split(",") # The output node names are used to select the usefull nodes
        ) 

        # Finally we serialize and dump the output graph to the filesystem
        with tf.gfile.GFile(output_graph, "wb") as f:
            f.write(output_graph_def.SerializeToString())
        print("%d ops in the final graph." % len(output_graph_def.node))

    return output_graph_def

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_dir", type=str, default="", help="Model folder to export")
    parser.add_argument("--output_node_names", type=str, default="", help="The name of the output nodes, comma separated.")
    args = parser.parse_args()

    freeze_graph(args.model_dir, args.output_node_names)