Model Freezing

Model Conversion and Storage with sess.run()

For TensorFlow training with sess.run, saver = tf.train.Saver() and saver.save() are used to save models. The following files are generated after each saver.save() call:

• checkpoint: lists the latest checkpoint file and other checkpoint files.
• model.ckpt.data-00000-of-00001: saves the current parameters, that is, the weight data.
• model.ckpt.index: saves the current parameter names.
• model.ckpt.meta: saves the current graph structure.

In this mode, the model weight data and model graph are saved separately. In the inference scenario, the TensorFlow freeze_graph function is used to combine the weight data and model graph into a .pb file, as shown in the dotted box in the following workflow diagram.

The process of generating a .pb file using the freeze_graph function of TensorFlow is as follows:

1. Specify the network model and checkpoint file path.
2. Define the input node. For example the input node for training is IteratorV2, but the input node required for inference is a placeholder.
3. Define the output node. For training, the output node is the node that produces the loss value. For inference, the required output is usually the node before the loss node, such as ArgMax or BiasAdd.
4. The same operator could be processed differently in the training graph and inference graph (for example, BatchNorm and Dropout operators). Therefore, you must use the model to generate an inference graph.
   • BatchNorm: During training, its mean and variance values are calculated based on the samples. However, during inference, they are calculated based on the moving average of the samples. Therefore, BatchNorm requires different mean computation modes for training and inference.
   • Dropout: During inference, you must mask dropout by setting rate to 1.

     1 2 3 4

     if is_training: x = npu_ops.dropout(x, 0.65) else: x = npu_ops.dropout(x, 1.0)

   Find the entry point function of the inference test logic in the training script and set is_training to False to generate an inference graph.

   1 2	# Call the network to generate an inference graph. alexnet.inference is the entry point function of the inference test logic in the training script. logits = alexnet.inference(inputs, version="he_uniform", num_classes=1000, is_training=False)

5. Call tf.train.writegraph to write the preceding inference graph to a .pb file, which will be fed into the freeze_graph function call.
6. Call freeze_graph to merge the .pb graph file generated by tf.train.writegraph and the checkpoint file to generate a .pb graph file for inference.

A code sample is provided as follows.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

import tensorflow as tf
from tensorflow.python.tools import freeze_graph
from npu_bridge.npu_init import *

# Import the model file.
import alexnet
# Specify the checkpoint path.
ckpt_path = "/opt/npu/model_ckpt/alexnet/model_8p/model.ckpt-0"

def main(): 
    tf.reset_default_graph()
    # Define the input node of the network.
    inputs = tf.placeholder(tf.float32, shape=[None, 224, 224, 3], name="input")
    # Call the network to generate an inference graph.
    logits = alexnet.inference(inputs, version="he_uniform",
                                  num_classes=1000, is_training=False)
    # Define the output node of the network.
    predict_class = tf.argmax(logits, axis=1, output_type=tf.int32, name="output")
    with tf.Session() as sess:
        # Save the graph to the model.pb file in the ./pb_model directory.
        # The model.pb file will be provided as input_graph to the following freeze_graph call.
        tf.train.write_graph(sess.graph_def, './pb_model', 'model.pb')    # Generate a model file.
        freeze_graph.freeze_graph(
		        input_graph='./pb_model/model.pb',   # Pass the model file generated using write_graph.
		        input_saver='',
		        input_binary=False, 
		        input_checkpoint=ckpt_path, # Pass the checkpoint file generated in training.
		        output_node_names='output',  # Consistent with the output node of the inference network.
		        restore_op_name='save/restore_all',
		        filename_tensor_name='save/Const:0',
		        output_graph='./pb_model/alexnet.pb',   # Set it to the name of the inference network to be generated.
		        clear_devices=False,
		        initializer_nodes='')
    print("done")

if __name__ == '__main__': 
    main()

Upon execution completion, the alexnet.pb file is generated in the ./pb_model/ directory. This file is the converted .pb graph file ready for inference.

For details about the dependent environment variables, see Training with a Single Device.

Model Conversion and Storage with Estimator

Estimator can save models in ckpt and saved_model formats. The ckpt method is similar to the sess.run method. You are advised to save the model in saved_model format to save the memory and avoid some possible errors. Generally, the saved_model model is saved through estimator.export_savedmodel, which consists of the following parts:

1
2
3
4

|--- saved_model.pb     # Network structure to be saved
|--- variables         # Parameter weight, including variables (tf.variable objects) of all models.
|---|--- variables.data-00000-of-00001
|---|--- variables.index

To convert saved_model into a .pb model for inference, perform the following steps:

1. Define the input node.
   The input received by Estimator during training is in Iterator format, which facilitates iteration between epochs. Before saving a model for inference, use the placeholder to define a specific input.

   1 2 3 4 5 6

   def serving_input_fn(): input_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_ids') input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({ 'input_ids': input_ids, }) return input_fn

2. Save the saved_model model.

   The Estimator model can save the model by directly calling the export_savedmodel function, which automatically switches modes and freezes the graph.

   1 2 3

   if FLAGS.do_export: estimator.evaluate() estimator.export_savedmodel(FLAGS.output_dir, serving_input_fn)

3. Freeze the .pb model.

   Use the freeze_graph function of TensorFlow to freeze a graph into a .pb model. Note that if your model has an NPU custom operator, import the NPU operator module to the source code of freeze_graph.

   1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17

   import tensorflow as tf from tensorflow.python.tools import freeze_graph from npu_bridge.npu_init import * freeze_graph.freeze_graph( input_saved_model_dir='savedModel', output_node_names='output', # Consistent with the output node of the inference network. output_graph='test.pb', # Name of the inference network to be generated initializer_nodes='', input_graph= None, input_saver= False, input_binary=False, input_checkpoint=None, restore_op_name=None, filename_tensor_name=None, clear_devices=False, input_meta_graph=False)

For details about the complete code example, see ModelZoo-TensorFlow.