Skip to content

TensorFlow Example 1*

MNIST Model Compilation Workflow*

MNIST is an entry-level computer vision dataset. Its input is handwritten digit images with a resolution of 28x28 pixels, and its output is the probability of the image corresponding to digits 0-9. Below, using the mnist model (TensorFlow v1.15) included with TensorFlow as an example, we illustrate the usage of the gxnpuc toolchain and API.

1. Generating NPU Files*

The MNIST computation model in this example is very simple and can be represented by a single formula: y = x * W + b (During training, softmax is also computed, but since we only need the index of the maximum value in the result during inference, and softmax is a monotonically increasing function, omitting this function does not affect the result). Here, x is the input data, y is the output data, and W and b are the trained parameters. The training process involves continuously adjusting W and b based on the computed y and the expected y_. On the NPU, we only need the trained W and b, not the training process.

1.1 Generating CKPT and PB Files*

First, we need to generate CKPT and PB files. Additionally, to conveniently specify the model's input and output nodes during NPU compilation, we can assign names to the input and output nodes. See the highlighted parts of the main function for details.

def main(_):
  # Import data
  mnist = input_data.read_data_sets(FLAGS.data_dir)

  # Create the model
  x = tf.placeholder(tf.float32, [None, 784], name="input_x") # Specify input name as input_x for easy use in compilation config scripts
  w = tf.Variable(tf.zeros([784, 10]))
  b = tf.Variable(tf.zeros([10]))
  y = tf.matmul(x, w) + b
  y = tf.identity(name="result") # Specify output name as result for easy use in compilation config scripts

  # Define loss and optimizer
  y_ = tf.placeholder(tf.int64, [None])

  # The raw formulation of cross-entropy,
  #
  #   tf.reduce_mean(-tf.reduce_sum(y_ * tf.math.log(tf.nn.softmax(y)),
  #                                 reduction_indices=[1]))
  #
  # can be numerically unstable.
  #
  # So here we use tf.compat.v1.losses.sparse_softmax_cross_entropy on the raw
  # logit outputs of 'y', and then average across the batch.
  cross_entropy = tf.losses.sparse_softmax_cross_entropy(labels=y_, logits=y)
  train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)

  config = tf.ConfigProto()
  jit_level = 0
  if FLAGS.xla:
    # Turns on XLA JIT compilation.
    jit_level = tf.OptimizerOptions.ON_1

  config.graph_options.optimizer_options.global_jit_level = jit_level
  run_metadata = tf.RunMetadata()
  sess = tf.compat.v1.Session(config=config)
  tf.global_variables_initializer().run(session=sess)
  # Train
  train_loops = 1000
  for i in range(train_loops):
    batch_xs, batch_ys = mnist.train.next_batch(100)

    # Create a timeline for the last loop and export to json to view with
    # chrome://tracing/.
    if i == train_loops - 1:
      sess.run(train_step,
               feed_dict={x: batch_xs,
                          y_: batch_ys},
               options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE),
               run_metadata=run_metadata)
      trace = timeline.Timeline(step_stats=run_metadata.step_stats)
      with open('/tmp/timeline.ctf.json', 'w') as trace_file:
        trace_file.write(trace.generate_chrome_trace_format())
    else:
      sess.run(train_step, feed_dict={x: batch_xs, y_: batch_ys})

  # Test trained model
  correct_prediction = tf.equal(tf.argmax(y, 1), y_)
  accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
  print(sess.run(accuracy,
                 feed_dict={x: mnist.test.images,
                            y_: mnist.test.labels}))

  # Generate CKPT and PB files
  saver = tf.train.Saver()
  saver.save(sess, "./mnist.ckpt")
  tf.train.write_graph(sess.graph_def, "./", "mnist.pb")

  sess.close()

After running the program, the mnist.ckpt.* and mnist.pb files will be generated in the current directory.

Note

Be sure to remember the model's input, output op names and their shapes, as they will be needed in the yaml configuration file for the gxnpuc compilation tool later.

1.2 Merging CKPT and PB Files into a FROZEN_PB File*

Use the freeze_graph.py script to merge mnist.ckpt.* and mnist.pb into a single pb file.

Note

The freeze_graph.py script may differ across TensorFlow versions.

Execute the command: 

$ python freeze_graph.py --input_graph=mnist.pb --input_checkpoint=./mnist.ckpt --output_graph=mnist_with_ckpt.pb --output_node_names=result
This generates the mnist_with_ckpt.pb file. Here, --input_graph is followed by the input PB name, --input_checkpoint is followed by the input CKPT name, --output_graph is followed by the merged FROZEN_PB file name, and --output_node_names is followed by the output node names (separated by commas if multiple). After execution, the mnist_with_ckpt.pb file is generated in the current directory.

If the model is saved using the saved_model format, use the following command to generate the FROZEN_PB file: 

$ python freeze_graph.py --input_saved_model_dir=./saved_model_dir --output_graph=mnist_with_ckpt.pb --output_node_names=result

1.3 Editing the NPU Configuration File*

Edit the configuration file mnist_config.yaml. The meanings are explained in the comments.

mnist_config.yaml
CORENAME: APUS # Chip model
MODEL_FILE: mnist_with_ckpt.pb # Input PB file
IN_FEATS_FILE: feats.txt # Input feature values, should cover model application scenarios as much as possible
QUANT_FILE: quant.yaml # Output quantization file
OUTPUT_FILE: mnist.h # Output NPU file name
COMPRESS: true # Compress fully connected layer weights
OUTPUT_TYPE: c_code # NPU file type
INPUT_OPS:
    input_x: [1, 784] # Input node name and data dimensions; each run inputs 1x784 data, i.e., one image
OUTPUT_OPS: [result] # Output node name
FUSE_BN: true # Merge BN parameters into convolution (if applicable)
MAX_CACHE_SIZE: 0 # Allocate CACHE memory for storing weights and data
USE_DATA_CACHE: false # Data is not stored in CACHE

Note

Here, input_x, result correspond one-to-one with the model's input, output op names.

1.4 Compilation*

First, use the gxnpuc tool to compile and generate the quantization file: 

$ gxnpuc mnist_config.yaml -q
Next, use the gxnpuc tool to compile and generate the NPU file mnist.h: 
$ gxnpuc mnist_config.yaml
The required memory information for the model is printed:

------------------------
Memory allocation info:
Mem1(data): 40
Mem2(instruction): 100
Mem3(in): 1568
Mem4(out): 20
Mem5(cache): 0
Mem6(weights): 15712
Total NPU Size (Mem0+Mem1+Mem2+Mem5+Mem6): 15852
Total Memory Size: 17440

The memory regions are described as follows:

Memory Region Description
Mem1(data) Intermediate data memory
Mem2(instruction) Instruction memory
Mem3(in) Input data memory
Mem4(out) Output data memory
Mem5(cache) Weights and data memory in SRAM
Mem6(weights) Weights memory

2. Executing NPU Files*

After the NPU file is generated, the model needs to be deployed and run on the GX830X development board.