Skip to content

Tensorflow model exporting note

kx9 edited this page Nov 22, 2017 · 2 revisions

uTensor

uTensor is an machine learning inferencing framework designed for resource constrained systems. At the introductory stage, the project focuses on running TensorFlow models on micro-controller systems. TensorFlow models are represented as a Graph by protocol buffer, therefore the node information has to be extracted and exported. Most TensorFlow models consist of floating point data. Quantizing of the models is essential to make inferencing practical on resource limited devices. TF-viewer is an intermediate tool used to weights in models to idx files. uTensor runtime loads these idx files and evaluate the graph defined in its source code.

Train the model in Tensorflow:

Same routine as in Tensorflow examples. For our example, we used the model described in tf-node-viewer.

Tensorflow model export

The model should be exported to a .pb file (protocol buffer file) which contains the weight and graph definition. Here is what it should look like at the end of your training script:

 saver.save(sess, “./my-model/model.ckpt”)
 out_nodes = [y_pred.op.name, y_.op.name, cross_entropy.op.name, correct_prediction.op.name, accuracy.op.name]
 sub_graph_def = gu.convert_variables_to_constants(sess, sess.graph_def, out_nodes)
graph_path = tf.train.write_graph(sub_graph_def, “./my-model”, “train.pb”, as_text=False)

At this point, train.pb should have been saved in your model directory.

Quantization

Use quantize_graph script to convert ./my-model/train.pb to an 8-bit quantized graph for computable in MCU device. Pete Warden

  1. download the script quantize_graph.py from tensorflow github repo
  2. create a directory named graph_out
  3. run python3 quantize_graph.py --input my-model/train.pb --mode=eightbit --output graph_out/quantized_graph.pb --output_node_names="Prediction/y_pred"

Export quantized model (quantized_graph.pb) to idx files:

As shown in tf-viewer readme, node_view.GraphInspector is a class to help developer to convert .pb file to idx file.

#loading the graph
graph = load_graph("./graph_out/quantized_graph.pb", name="")
#setup the inspector with feed_dict
#feed indicates that the first image in the test-set is used to evaluate all nodes
inspector = GraphInspector(graph, feed_dict={"Placeholder:0": mnist.test.images[0:1]})
#saving the nodes
inspector.snap("node name you want to save")
#viewing node information

ls method would list the operators in the graph and snap method is convert node information to idx format. Running snap method would create a named by the corresponding op name by default. Two sub-directories corresponding to a node’s inputs and outputs are also created. These sub-directories contain the weights saved in idx format.

Move the idx files to SD card:

uTensor supports loading idx file for graph evaluation. You will need an SD card formatted to FAT32 (and less than 32GB). The exported idx files should be saved into the SD card. The file paths of individual idx files should reflect to where they are saved in the SD card.

Load weight from SD Card

In the uTensor framework, idx files can be loaded into tensors by the following:

Tensor<unsigned char> t = t_import.ubyte_import(“/fs/testData/idxImport/uint8_4d_power2.idx”, "t")

In the near future, we aim to streamline this process to make it more user friendly.