Guides
rknn
Running TFLite models with RKNN2

Introduction

RKNN2 is Rockchip's software suite to utilize the NPU on their new device platforms like the RK3588 and RK3566.

The NPU doesn't directly run models as is, you need to convert them into the custom .rknn format.

To do so, there is the RKNN-Toolkit2 which lets you analyze, quantize, convert and see how the model would run on the NPU layer wise.

Prequisites

This guide is a look into how you can convert and run some basic TFLite models, it targets the following platforms

  • RK3566/RK3568 (1.0 TOPs NPU)
  • RK3588/RK3588S (6.0 TOPs NPU)

You can check out https://github.com/sravansenthiln1/rknn_tflite (opens in a new tab) for the example code and other details.

Running an example model

Here's an example running the sample CNN model that predicts digits from the popular MNIST dataset.

mnist digit

The sample resides at https://github.com/sravansenthiln1/rknn_tflite/tree/main/digit_recognize (opens in a new tab)

and has the pre-trained and converted model to test.

Here's the run through of the code running the inference

Importing our essential libraries

import numpy as np
from PIL import Image
import time
 
from rknnlite.api import RKNNLite

Defining the file paths

we have to define the file paths for things such as the model and input.

Set the path to the model

MODEL_PATH = "./digit_recognize_28.rknn"

Set the path to the input image

IMAGE_PATH = "./digit7.png"

Creating a RKNNlite handle

Create a class instance for interacting with RKNN

rknn_lite = RKNNLite()

Load the RKNN model

Load the RKNN model using the previously created RKNNLite instance

ret = rknn_lite.load_rknn(MODEL_PATH)
if ret != 0:
    print('Load RKNN model failed')
    exit(ret)
print('done')

Load the image Object

Load the input image specified with the path we set

img = Image.open(IMAGE_PATH)
img = img.convert("L")
img = np.expand_dims(img, 0)

Initialize RKNNLite

Start up the NPU runtime libraries

print('--> Init runtime environment')
ret = rknn_lite.init_runtime()
if ret != 0:
    print('Init runtime environment failed!')
    exit(ret)
print('done')

Add extra outer dimension

RKNN supports batch processing, so we have to encapsulate the image object within another dimension

img = np.array([img])

Invoke the inference

Invoke the inference call and process all the input samples

st = time.time()
output = rknn_lite.inference(inputs=[img])
en = time.time()

Print the inferenced output and some time statistics

The time statistic is simply a start stop time difference, while the inference prediction is the index of the highest prediction confidence.

print("Inference in: ", (en - st) * 1000, "ms" )
print("predicts: ", np.argmax(output))

Remember the model, and image files should be in the same path as this script.

Example run

here is an example running on the Khadas Edge2 (opens in a new tab),

--> Load RKNN model
done
--> Init runtime environment
I RKNN: [08:25:37.584] RKNN Runtime Information: librknnrt version: 1.5.2 (c6b7b351a@2023-08-23T15:28:22)
I RKNN: [08:25:37.584] RKNN Driver Information: version: 0.7.2
W RKNN: [08:25:37.584] Current driver version: 0.7.2, recommend to upgrade the driver to the new version: >= 0.8.8
I RKNN: [08:25:37.584] RKNN Model Information: version: 6, toolkit version: 1.5.2+b642f30c(compiler version: 1.5.2 (c6b7b351a@2023-08-23T07:39:01)), target: RKNPU v2, target platform: rk3588, framework name: TFLite, framework layout: NHWC, model inference type: static_shape
done
Inference in:  0.6785392761230469 ms
predicts:  7

Notes:

To do any optimization on the model with the conversion step, you'll need to specify a sample input in the form of an sample image .png, or .jpg or .npy file.