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.