Preparing a model for OpenCV

TensorFlow generates multiple files when we train a new model, generating files for events that store accuracy and loss, and other metrics obtained in each step; also, some files will store the variable results obtained for each step or checkpoint. These variables are the weights that networks learn in training. But it's not convenient to share all these files in production, as OpenCV would not be able to manage them. At the same time, there are nodes that are used only for training and not for inference. We have to remove these nodes from the model, nodes such as dropouts layers or training input iterators.

To put our model into production, we need to do the following:

Freezing takes the graph definition and a set of checkpoints and merges them together into a single file, converting the variables into constants. To freeze our model, we have to move into the saved model folder and execute the following script provided by TensorFlow:

freeze_graph --input_graph=graph.pbtxt --input_checkpoint=model.ckpt-20000 --output_graph frozen_graph.pb --output_node_names=softmax_tensor

Now, we generate a new file called frozen_graph.pb, which is the merged and frozen graph. Then, we have to remove the input layers used for training purposes. If we review the graph using TensorBoard, we can see that our input to the first convolutional neural network is the IteratorGetNext node, which we have to cut and set as a single layer input of a 20 x 20 pixel image of one channel. Then, we can use the TensorFlow transform_graph application, which allows us to change the graph, cutting or modifying the TensorFlow model graph. To remove the layer connected to the ConvNet, we execute the following code:

transform_graph --in_graph="frozen_graph.pb" --out_graph="frozen_cut_graph.pb" --inputs="IteratorGetNext" --outputs="softmax_tensor" --transforms='strip_unused_nodes(type=half, shape="1,20,20,1") fold_constants(ignore_errors=true) fold_batch_norms fold_old_batch_norms sort_by_execution_order'
It's very important to add the sort_by_execution_order parameter to ensure that the layers are stored in order in the model graph, to allow OpenCV to correctly import the model. OpenCV sequentially imports the layers from the graph model, checking that all previous layers, operations, or variables are imported; if not, we will receive an import error. TensorEditor doesn't take care of the execution order in the graph to construct and execute it.

After executing transform_graph, we have a new model saved as frozen_cut_graph.pb. The final step requires us to optimize the graph, removing all training operations and layers such as dropout. We are going to use the following command to optimize our model for production/inference; this application is provided by TensorFlow:

optimize_for_inference.py --input frozen_cut_graph.pb --output frozen_cut_graph_opt.pb --frozen_graph True --input_names IteratorGetNext --output_names softmax_tensor

The output of this is a file called frozen_cut_graph_opt.pb. This file is our final model, which we can import and use in our OpenCV code.