Platform scope and source
This article follows the RKNN V2.5.0 vendor PDF. Only the chips and interfaces listed in the source are in scope. For an ArmSoM board, check the SoC, NPU driver, Toolkit/Runtime, SDK/BSP and kernel versions first. No ArmSoM hardware validation has been performed; source examples and measurements are not ArmSoM test results.
Series navigation
- Requirements and object lifecycle
- Model configuration and quantization parameters
- Loading Caffe, TensorFlow, ONNX and PyTorch
- Building, hybrid quantization and export
- Runtime, operator registration and inference
- Accuracy, performance, memory and device queries
- Generating, building and running C++ deployment examples
4 Model Loading APIs
RKNN-Toolkit2 currently supports load non-RKNN models of Caffe, TensorFlow, TensorFlow Lite, ONNX, DarkNet, PyTorch. There are different calling interfaces when loading models, the loading interfaces are described in detail below.
4.1 Loading Caffe model
| API | load_caffe |
|---|---|
| Description | Load Caffe model. (Unavailable in ARM64) |
| Parameter | model: The path of Caffe model structure file (suffixed with “.prototxt” ). |
| Parameter | blobs: The path of Caffe model binary data file (suffixed with “.caffemodel”). |
| Parameter | input_name: When the Caffe model has multiple inputs, you can specify the order of the input layer names through this parameter, such as [‘input1’,’input2’,’input3’],note that the name needs to be consistent with the model input name;The default value is None, means the sequence is automatically given by the Caffe model file (file suffix with .prototxt). |
| Return Value |
0: Import successfully. |
| Return Value |
-1: Import failed. |
The sample code is as follows:
# Load the mobilenet_v2 Caffe model in the current path
ret = rknn.load_caffe(model='./mobilenet_v2.prototxt',
blobs='./mobilenet_v2.caffemodel')
4.2 Loading TensorFlow model
| API | load_tensorflow |
|---|---|
| Description | Load TensorFlow model. (Unavailable in ARM64) |
| Parameter | tf_pb: The path of TensorFlow model file (suffixed with “.pb”). |
| Parameter | inputs: The input node (operand name) of model, input with multiple nodes is supported now. All the input node string are placed in a list. |
| Parameter | input_size_list: The shapes of input node, all the input shape are placed in a list. As in the example of ssd_mobilenet_v1 model, the input_size_list parameter should be set to [[1,300,300,3]]. |
| Parameter | outputs: The output node (operand name) of model, output with multiple nodes is supported now. All the output nodes are placed in a list. |
| Parameter | input_is_nchw: Whether the input layout of the model is already NCHW. The default value is False, means the default input layout is NHWC. |
| Return value |
0: Import successfully. |
| Return value |
-1: Import failed. |
The sample code is as follows:
# Load ssd_mobilenet_v1_coco_2017_11_17 TF model in the current path
ret = rknn.load_tensorflow(tf_pb='./ssd_mobilenet_v1_coco_2017_11_17.pb',
inputs=['Preprocessor/sub'],
outputs=['concat', 'concat_1'],
input_size_list=[[300, 300, 3]])
4.3 Loading TensorFlow Lite model
| API | load_tflite |
|---|---|
| Description | Load TensorFlow Lite model. (Unavailable in ARM64) |
| Parameter | model: The path of TensorFlow Lite model file (suffixed with “.tflite”). |
| Parameter | input_is_nchw: Whether the input layout of the model is already NCHW. The default value is False, that is, the default input layout is NHWC. |
| Return Value |
0: Import successfully. |
| Return Value |
-1: Import failed. |
The sample code is as follows:
# Load the mobilenet_v1 TF-Lite model in the current path
ret = rknn.load_tflite(model='./mobilenet_v1.tflite')
4.4 Loading ONNX model
| API | load_onnx |
|---|---|
| Description | Load ONNX model. |
| Parameter | model: The path of ONNX model file (suffixed with “.onnx”). |
| Parameter | inputs: The input node (operand name) of model, input with multiple nodes is supported now. All the input node string are placed in a list. The default value is None, means get from model. |
| Parameter | input_size_list: The shapes of input node, all the input shape are placed in a list. If inputs set, the input_size_list should be set also. defualt is None. |
| Parameter | input_initial_val: Set the initial value of the model input, the format is ndarray list. The default value is None. Mainly used to fix some input as constant, For the input that does not need to be fix as a constant, it can be set to None, such as [None, np.array([1])]. |
| Parameter | outputs: The output node (operand name) of model, output with multiple nodes is supported now. All the output nodes are placed in a list. The default value is None, means get from model. |
| Return Value |
0: Import successfully. |
| Return Value |
-1: Import failed. |
The sample code is as follows:
# Load the arcface onnx model in the current path
ret = rknn.load_onnx(model='./arcface.onnx')
4.5 Loading DarkNet model
| API | load_darknet |
|---|---|
| Description | Load DarkNet model. (Unavailable in ARM64) |
| Parameter | model: The path of DarkNet model structure file (suffixed with “.cfg”). |
| Parameter | weight: The path of weight file (suffixed with “.weight”). |
| Return Value | 0: Import successfully. |
| Return Value | -1: Import failed. |
The sample code is as follows:
# Load the yolov3-tiny DarkNet model in the current path
ret = rknn.load_darknet(model='./yolov3-tiny.cfg',
weight= './yolov3.weights')
4.6 Loading PyTorch model
| API | load_pytorch |
|---|---|
| Description | Load PyTorch model. Support the Quantization Aware Training (QAT) model, but need update torch version to ‘1.9.0’ or higher. |
| Parameter | model: The path of PyTorch model structure file (suffixed with “.pt”), and need a model in the torchscript format. |
| Parameter | input_size_list: The shapes of input node, all the input shape are placed in a list. |
| Return Value |
0: Import successfully. |
| Return Value |
-1: Import failed. |
The sample code is as follows:
# Load the PyTorch model resnet18 in the current path
ret = rknn.load_pytorch(model='./resnet18.pt',
input_size_list=[[1,3,224,224]])
Source and editorial notes
Technical chapters are split by use. Code, identifiers and parameters are retained. Chinese follows the Chinese source; English, Russian and German follow the English source. No invented figures are added to chapters without figures.
Editorial notes: Example logs still show SDK 1.5.2 and remain historical examples, not V2.5.0 measurements. The Ubuntu20.04 Python cell differs (Chinese 3.8/3.9; English 3.8/3.8, also differing for Debian11); the ARM64 section still labels Debian x64. Source values are retained; verify the actual installation package. V2.5.0 adds the C++ codegen chapter and refines quantization and memory strategies; older parameters have not replaced the new source.