Quantizing a Model for the Rockchip NPU — MNIST RKNN Walkthrough

Why this walkthrough still matters

Quantization is the single biggest lever for running a neural network on a Rockchip NPU: it shrinks the model and lets the accelerator use its INT8 datapath. The original MNIST example is small, but the shape of the pipeline — train → freeze → convert → calibrate/quantize → verify — is exactly what you repeat for every real model on RK3399Pro, RK3566/68, RK3588, and RK3576.

1. Build the network (model.py)

A LeNet-5-style network. The input placeholder is shaped [None, 28, 28] — a single-channel 28×28 image flattened in the batch — which keeps the later RKNN quantization step straightforward.

import tensorflow as tf

# input uses 28x28 layout to make later rknn quantization easier
x = tf.placeholder("float", [None, 28, 28], name='x')
y_ = tf.placeholder("float", [None, 10], name='y_')
keep_prob = tf.placeholder("float", name='keep_prob')

def weight_variable(shape, name):
    initial = tf.truncated_normal(shape, stddev=0.1)
    return tf.Variable(initial, name=name)

def bias_variable(shape, name):
    initial = tf.constant(0.1, shape=shape)
    return tf.Variable(initial, name=name)

def conv2d(x, W):
    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')

def max_pool_2x2(x):
    return tf.nn.max_pool(x, ksize=[1, 2, 2, 1],
                          strides=[1, 2, 2, 1], padding='SAME')

# convolution layer
def lenet5_layer(input, weight, bias, weight_name, bias_name):
    W_conv = weight_variable(weight, weight_name)
    b_conv = bias_variable(bias, bias_name)
    h_conv = tf.nn.relu(conv2d(input, W_conv) + b_conv)
    return max_pool_2x2(h_conv)

# fully-connected layer
def dense_layer(layer, weight, bias, weight_name, bias_name):
    W_fc = weight_variable(weight, weight_name)
    b_fc = bias_variable(bias, bias_name)
    return tf.nn.relu(tf.matmul(layer, W_fc) + b_fc)

def build_model(is_training):
    # first conv
    x_image = tf.reshape(x, [-1, 28, 28, 1])
    W_conv1 = [5, 5, 1, 32]
    b_conv1 = [32]
    layer1 = lenet5_layer(x_image, W_conv1, b_conv1, 'W_conv1', 'b_conv1')
    # second conv
    W_conv2 = [5, 5, 32, 64]
    b_conv2 = [64]
    layer2 = lenet5_layer(layer1, W_conv2, b_conv2, 'W_conv2', 'b_conv2')
    # fully-connected
    W_fc1 = [7 * 7 * 64, 1024]
    b_fc1 = [1024]
    layer2_flat = tf.reshape(layer2, [-1, 7*7*64])
    layer3 = dense_layer(layer2_flat, W_fc1, b_fc1, 'W_fc1', 'b_fc1')
    # softmax
    W_fc2 = weight_variable([1024, 10], 'W_fc2')
    b_fc2 = bias_variable([10], 'b_fc2')
    if is_training:
        # dropout only during training
        h_fc1_drop = tf.nn.dropout(layer3, keep_prob)
        finaloutput = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2, name="y_conv")
    else:
        finaloutput = tf.nn.softmax(tf.matmul(layer3, W_fc2) + b_fc2, name="y_conv")
    print('finaloutput:', finaloutput)
    return finaloutput

2. Train the network (train.py)

A create_training_graph(is_quantify) switch lets you optionally insert TensorFlow's fake-quantization training graph. The base run uses is_quantify=False. Note the reshape_batch helper that turns MNIST's flat 784-vector into a 28×28 image.

# -*- coding=utf-8 -*-
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
from model import build_model, x, y_, keep_prob

mnist = input_data.read_data_sets("../MNIST_data/", one_hot=True)

def create_training_graph(is_quantify):
    # build the training graph; create_training_graph is inserted here
    g = tf.get_default_graph()  # arg for create_training_graph (default graph)
    # call the network definition to get the output
    y_conv = build_model(is_training=True)  # is_training=True (dropout used)
    # loss function
    cross_entropy = -tf.reduce_sum(y_ * tf.log(y_conv))
    print('cost:', cross_entropy)
    if is_quantify:
        # insert create_training_graph AFTER loss, BEFORE optimize
        tf.contrib.quantize.create_training_graph(input_graph=g, quant_delay=0)
    # optimize
    optimize = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)
    # accuracy
    correct_prediction = tf.equal(tf.argmax(y_conv, 1), tf.argmax(y_, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
    return dict(
        x=x,
        y=y_,
        keep_prob=keep_prob,
        optimize=optimize,
        cost=cross_entropy,
        correct_prediction=correct_prediction,
        accuracy=accuracy,
    )

def reshape_batch(batch):
    rebatch = []
    for item in batch:
        b = item.reshape(28, 28)
        rebatch.append(b)
    return rebatch

def train_network(graph, ckpt, point_dir, pbtxt):
    init = tf.global_variables_initializer()
    saver = tf.train.Saver()
    with tf.Session() as sess:
        sess.run(init)
        # train 20000 steps; accuracy reaches 99%+
        for i in range(20000):
            batch = mnist.train.next_batch(50)
            if i % 100 == 0:
                train_accuracy = sess.run([graph['accuracy']], feed_dict={
                    graph['x']: reshape_batch(batch[0]),  # image data
                    graph['y']: batch[1],                  # labels
                    graph['keep_prob']: 1.0})
                print("step %d, training accuracy %g" % (i, train_accuracy[0]))
            sess.run([graph['optimize']], feed_dict={
                graph['x']: reshape_batch(batch[0]),
                graph['y']: batch[1],
                graph['keep_prob']: 0.5})
        test_accuracy = sess.run([graph['accuracy']], feed_dict={
            graph['x']: reshape_batch(mnist.test.images),
            graph['y']: mnist.test.labels,
            graph['keep_prob']: 1.0})
        print("Test accuracy %g" % test_accuracy[0])
        # save ckpt and pbtxt (change paths to your own)
        saver.save(sess, ckpt)
        tf.train.write_graph(sess.graph_def, point_dir, pbtxt, True)
        print(tf.trainable_variables())
        print(tf.get_variable('W_fc2', [1024, 10]).value)

if __name__ == "__main__":
    ckpt = './checkpoint/mnist.ckpt'
    point_dir = './checkpoint'
    pbtxt = 'mnist.pbtxt'
    g1 = create_training_graph(False)
    train_network(g1, ckpt, point_dir, pbtxt)

3. Freeze to a PB graph (freese.py)

Read the checkpoint and write a frozen PB — all variables become inline constants, and only the y_conv output node is kept. With is_quantify=False, the eval (fake-quant) graph is not inserted.

import tensorflow as tf
import os.path
from model import build_model
from tensorflow.python.framework import graph_util

def create_inference_graph():
    # build the model for evaluation
    logits = build_model(is_training=False)
    return logits

def load_variables_from_checkpoint(sess, start_checkpoint):
    # restore weights from the checkpoint
    saver = tf.train.Saver(tf.global_variables())
    saver.restore(sess, start_checkpoint)

def frozen(is_quantify, ckpt, pbtxt):
    init = tf.global_variables_initializer()
    with tf.Session() as sess:
        sess.run(init)
        logits = create_inference_graph()
        if is_quantify:
            tf.contrib.quantize.create_eval_graph()
        load_variables_from_checkpoint(sess, ckpt)
        # turn all variables into inline constants and save
        frozen_graph_def = graph_util.convert_variables_to_constants(
            sess, sess.graph_def, ['y_conv'])
        tf.train.write_graph(
            frozen_graph_def,
            os.path.dirname(pbtxt),
            os.path.basename(pbtxt),
            as_text=False)
        tf.logging.info('Saved frozen graph to %s', pbtxt)

if __name__ == "__main__":
    ckpt = './checkpoint/mnist.ckpt'
    pbtxt = 'mnist_frozen_graph.pb'
    frozen(False, ckpt, pbtxt)

4. Convert PB → RKNN (with optional quantization)

First extract the test-set images and a dataset.txt (needed for quantization calibration). Then run the conversion: common_transfer builds an unquantized RKNN, and quantify_transfer builds a quantized one using quantized_dtype='dynamic_fixed_point-8'.

4.1 Generate images + dataset.txt (get_image.py)

import struct
import numpy as np
import PIL.Image
from PIL import Image
import os

os.system("mkdir ../MNIST_data/mnist_test")
filename = '../MNIST_data/t10k-images.idx3-ubyte'
dataset = './dataset.txt'
binfile = open(filename, 'rb')
buf = binfile.read()
index = 0
data_list = []
magic, numImages, numRows, numColumns = struct.unpack_from('>IIII', buf, index)
index += struct.calcsize('>IIII')
for image in range(0, numImages):
    im = struct.unpack_from('>784B', buf, index)
    index += struct.calcsize('>784B')
    im = np.array(im, dtype='uint8')
    im = im.reshape(28, 28)
    im = Image.fromarray(im)
    im.save('../MNIST_data/mnist_test/test_%s.jpg' % image, 'jpeg')
    data_list.append('../MNIST_data/mnist_test/test_%s.jpg\n' % image)
with open(dataset, 'w+') as ff:
    ff.writelines(data_list)

4.2 PB → RKNN conversion (rknn_transfer.py)

from rknn.api import RKNN

def common_transfer(pb_name, export_name):
    ret = 0
    rknn = RKNN()
    print('--> Loading model')
    ret = rknn.load_tensorflow(
        tf_pb='./mnist_frozen_graph.pb',
        inputs=['x'],
        outputs=['y_conv'],
        input_size_list=[[28, 28, 1]])
    if ret != 0:
        print('load_tensorflow error')
        rknn.release()
        return ret
    print('done')
    print('--> Building model')
    rknn.build(do_quantization=False)
    print('done')
    rknn.export_rknn('./mnist.rknn')
    rknn.release()
    return ret

def quantify_transfer(pb_name, dataset_name, export_name):
    ret = 0
    print(pb_name, dataset_name, export_name)
    rknn = RKNN()
    rknn.config(channel_mean_value='', reorder_channel='',
                quantized_dtype='dynamic_fixed_point-8')
    print('--> Loading model')
    ret = rknn.load_tensorflow(
        tf_pb=pb_name,
        inputs=['x'],
        outputs=['y_conv'],
        input_size_list=[[28, 28, 1]])
    if ret != 0:
        print('load_tensorflow error')
        rknn.release()
        return ret
    print('done')
    print('--> Building model')
    rknn.build(do_quantization=True, dataset=dataset_name)
    print('done')
    rknn.export_rknn(export_name)
    rknn.release()
    return ret

if __name__ == '__main__':
    pb_name = './mnist_frozen_graph.pb'
    export_name = './mnist.rknn'
    ret = common_transfer(pb_name, export_name)
    if ret != 0:
        print('======common transfer error !!===========')
    else:
        print('======common transfer ok !!===========')
    dataset_name = './dataset.txt'
    export_name = './mnist_quantization.rknn'
    quantify_transfer(pb_name, dataset_name, export_name)
    if ret != 0:
        print('======quantization transfer 10000 error !!===========')
    else:
        print('======quantization transfer 10000 ok !!===========')

5. Compare PB vs RKNN inference

Run the original frozen-graph inference and the RKNN inference over the same images and compare accuracy. The RKNN side loads the model, calls init_runtime() (this talks to the NPU), then inference().

# rknn_predict.py (key parts)
import numpy as np
from PIL import Image
from rknn.api import RKNN

def get_predict(probability):
    data = probability[0][0]
    data = data.tolist()
    max_prob = max(data)
    return data.index(max_prob), max_prob

def load_model(model_name):
    rknn = RKNN()
    print('-->loading model')
    rknn.load_rknn(model_name)
    print('loading model done')
    print('--> Init runtime environment')
    ret = rknn.init_runtime()
    if ret != 0:
        print('Init runtime environment failed')
        exit(ret)
    print('done')
    return rknn

def predict(rknn, length):
    acc_count = 0
    for i in range(length):
        im = Image.open("../MNIST_data/mnist_test/test_%d.jpg" % i)
        im = im.resize((28, 28), Image.ANTIALIAS)
        im = np.asarray(im)
        outputs = rknn.inference(inputs=[im])
        pred, prob = get_predict(outputs)
        if list(mnist.test.labels[i]).index(1) == pred:
            acc_count += 1
    result = float(acc_count) / length
    print('result:', result)

if __name__ == "__main__":
    # switch to the quantized or unquantized rknn model as needed
    model_name = './mnist.rknn'
    length = 10000
    rknn = load_model(model_name)
    predict(rknn, length)
    rknn.release()

The original thread includes an accuracy-comparison table (visible to logged-in forum users) and reports that the quantized RKNN model's file size is roughly half of the unquantized RKNN, which in turn is roughly half of the float32 PB — consistent with float32 → int16 → int8.

6. Going further: fake quantization (optional)

The thread's appendix shows a full fake-quantization path: re-train with tf.contrib.quantize.create_training_graph / create_eval_graph, freeze, convert the fake-quantized PB to a fully-quantized TFLite via toco, then load that TFLite into RKNN. This generally gives better INT8 accuracy than post-training quantization alone because the network learns to tolerate the reduced precision during training.

# toco: pb -> fully-quantized tflite (int8 in/out/weights)
toco \
 --graph_def_file=mnist_frozen_graph_fake.pb \
 --output_file=mnist_fakequantize.tflite \
 --output_format=TFLITE \
 --inference_type=QUANTIZED_UINT8 \
 --input_shapes=1,28,28 \
 --input_arrays=x \
 --output_arrays=y_conv \
 --mean_values=0 \
 --std_dev_values=256 \
 --change_concat_input_ranges=false --allow_custom_ops

Then rknn.load_tflite(model='./mnist_fakequantize.tflite') and rknn.build(do_quantization=False) to export the RKNN. The thread closes with a community Q&A covering multi-input handling (later RKNN Toolkit versions support it natively), out-of-memory during training, board-vs-PC conversion size differences, and a shape-mismatch pitfall when the placeholder x is defined as [None,28,28] but RKNN adds the channel dimension per input_size_list.

Recap

• Train → freeze → convert → quantize → verify is the invariant NPU pipeline, regardless of SoC.
• The calibration dataset.txt drives quantization quality — make it representative.
• Verify NPU output against the original float model before you trust the INT8 result.
• Original author: peng (Rockchip / Toybrick community). Source: t.rock-chips.com/forum.php?mod=viewthread&tid=858.