1 // Copyright (c) 2021 by Rockchip Electronics Co., Ltd. All Rights Reserved.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 /*-------------------------------------------
16 Includes
17 -------------------------------------------*/
18 #include "rk_mpi_mmz.h"
19 #include "rknn_api.h"
20
21 #include <float.h>
22 #include <stdio.h>
23 #include <stdlib.h>
24 #include <string.h>
25 #include <sys/time.h>
26
27 #define STB_IMAGE_IMPLEMENTATION
28 #include "stb/stb_image.h"
29 #define STB_IMAGE_RESIZE_IMPLEMENTATION
30 #include <stb/stb_image_resize.h>
31
32 /*-------------------------------------------
33 Functions
34 -------------------------------------------*/
getCurrentTimeUs()35 static inline int64_t getCurrentTimeUs()
36 {
37 struct timeval tv;
38 gettimeofday(&tv, NULL);
39 return tv.tv_sec * 1000000 + tv.tv_usec;
40 }
41
rknn_GetTopN(float * pfProb,float * pfMaxProb,uint32_t * pMaxClass,uint32_t outputCount,uint32_t topNum)42 static int rknn_GetTopN(float* pfProb, float* pfMaxProb, uint32_t* pMaxClass, uint32_t outputCount, uint32_t topNum)
43 {
44 uint32_t i, j;
45 uint32_t top_count = outputCount > topNum ? topNum : outputCount;
46
47 for (i = 0; i < topNum; ++i) {
48 pfMaxProb[i] = -FLT_MAX;
49 pMaxClass[i] = -1;
50 }
51
52 for (j = 0; j < top_count; j++) {
53 for (i = 0; i < outputCount; i++) {
54 if ((i == *(pMaxClass + 0)) || (i == *(pMaxClass + 1)) || (i == *(pMaxClass + 2)) || (i == *(pMaxClass + 3)) ||
55 (i == *(pMaxClass + 4))) {
56 continue;
57 }
58
59 if (pfProb[i] > *(pfMaxProb + j)) {
60 *(pfMaxProb + j) = pfProb[i];
61 *(pMaxClass + j) = i;
62 }
63 }
64 }
65
66 return 1;
67 }
68
dump_tensor_attr(rknn_tensor_attr * attr)69 static void dump_tensor_attr(rknn_tensor_attr* attr)
70 {
71 printf(" index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
72 "zp=%d, scale=%f\n",
73 attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
74 attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
75 get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
76 }
77
load_image(const char * image_path,rknn_tensor_attr * input_attr)78 static unsigned char* load_image(const char* image_path, rknn_tensor_attr* input_attr)
79 {
80 int req_height = 0;
81 int req_width = 0;
82 int req_channel = 0;
83
84 switch (input_attr->fmt) {
85 case RKNN_TENSOR_NHWC:
86 req_height = input_attr->dims[1];
87 req_width = input_attr->dims[2];
88 req_channel = input_attr->dims[3];
89 break;
90 case RKNN_TENSOR_NCHW:
91 req_height = input_attr->dims[2];
92 req_width = input_attr->dims[3];
93 req_channel = input_attr->dims[1];
94 break;
95 default:
96 printf("meet unsupported layout\n");
97 return NULL;
98 }
99
100 int height = 0;
101 int width = 0;
102 int channel = 0;
103
104 unsigned char* image_data = stbi_load(image_path, &width, &height, &channel, req_channel);
105 if (image_data == NULL) {
106 printf("load image failed!\n");
107 return NULL;
108 }
109
110 if (width != req_width || height != req_height) {
111 unsigned char* image_resized = (unsigned char*)STBI_MALLOC(req_width * req_height * req_channel);
112 if (!image_resized) {
113 printf("malloc image failed!\n");
114 STBI_FREE(image_data);
115 return NULL;
116 }
117 if (stbir_resize_uint8(image_data, width, height, 0, image_resized, req_width, req_height, 0, channel) != 1) {
118 printf("resize image failed!\n");
119 STBI_FREE(image_data);
120 return NULL;
121 }
122 STBI_FREE(image_data);
123 image_data = image_resized;
124 }
125
126 return image_data;
127 }
128
129 /*-------------------------------------------
130 Main Functions
131 -------------------------------------------*/
main(int argc,char * argv[])132 int main(int argc, char* argv[])
133 {
134 if (argc < 3) {
135 printf("Usage:%s model_path input_path [loop_count]\n", argv[0]);
136 return -1;
137 }
138
139 char* model_path = argv[1];
140 char* input_path = argv[2];
141
142 int loop_count = 1;
143 if (argc > 3) {
144 loop_count = atoi(argv[3]);
145 }
146
147 rknn_context ctx = 0;
148
149 // Load RKNN Model
150 int ret = rknn_init(&ctx, model_path, 0, 0, NULL);
151 if (ret < 0) {
152 printf("rknn_init fail! ret=%d\n", ret);
153 return -1;
154 }
155
156 // Get sdk and driver version
157 rknn_sdk_version sdk_ver;
158 ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &sdk_ver, sizeof(sdk_ver));
159 if (ret != RKNN_SUCC) {
160 printf("rknn_query fail! ret=%d\n", ret);
161 return -1;
162 }
163 printf("rknn_api/rknnrt version: %s, driver version: %s\n", sdk_ver.api_version, sdk_ver.drv_version);
164
165 // Get Model Input Output Info
166 rknn_input_output_num io_num;
167 ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
168 if (ret != RKNN_SUCC) {
169 printf("rknn_query fail! ret=%d\n", ret);
170 return -1;
171 }
172 printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);
173
174 printf("input tensors:\n");
175 rknn_tensor_attr input_attrs[io_num.n_input];
176 memset(input_attrs, 0, io_num.n_input * sizeof(rknn_tensor_attr));
177 for (uint32_t i = 0; i < io_num.n_input; i++) {
178 input_attrs[i].index = i;
179 // query info
180 ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
181 if (ret < 0) {
182 printf("rknn_init error! ret=%d\n", ret);
183 return -1;
184 }
185 dump_tensor_attr(&input_attrs[i]);
186 }
187
188 printf("output tensors:\n");
189 rknn_tensor_attr output_attrs[io_num.n_output];
190 memset(output_attrs, 0, io_num.n_output * sizeof(rknn_tensor_attr));
191 for (uint32_t i = 0; i < io_num.n_output; i++) {
192 output_attrs[i].index = i;
193 // query info
194 ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
195 if (ret != RKNN_SUCC) {
196 printf("rknn_query fail! ret=%d\n", ret);
197 return -1;
198 }
199 dump_tensor_attr(&output_attrs[i]);
200 }
201
202 // Get custom string
203 rknn_custom_string custom_string;
204 ret = rknn_query(ctx, RKNN_QUERY_CUSTOM_STRING, &custom_string, sizeof(custom_string));
205 if (ret != RKNN_SUCC) {
206 printf("rknn_query fail! ret=%d\n", ret);
207 return -1;
208 }
209 printf("custom string: %s\n", custom_string.string);
210
211 unsigned char* input_data = NULL;
212 rknn_tensor_type input_type = RKNN_TENSOR_UINT8;
213 rknn_tensor_format input_layout = RKNN_TENSOR_NHWC;
214
215 // Load image
216 input_data = load_image(input_path, &input_attrs[0]);
217
218 if (!input_data) {
219 return -1;
220 }
221
222 int mb_flags = RK_MMZ_ALLOC_TYPE_CMA | RK_MMZ_ALLOC_UNCACHEABLE;
223
224 // Allocate input memory in outside
225 MB_BLK input_mb;
226 int input_size = input_attrs[0].size_with_stride;
227 ret = RK_MPI_MMZ_Alloc(&input_mb, input_size, mb_flags);
228 if (ret < 0) {
229 printf("RK_MPI_MMZ_Alloc failed, ret: %d\n", ret);
230 return ret;
231 }
232 void* input_virt = RK_MPI_MMZ_Handle2VirAddr(input_mb);
233 if (input_virt == NULL) {
234 printf("RK_MPI_MMZ_Handle2VirAddr failed!\n");
235 return -1;
236 }
237 uint64_t input_phys = RK_MPI_MMZ_Handle2PhysAddr(input_mb);
238 if (input_phys == 0) {
239 printf("RK_MPI_MMZ_Handle2PhysAddr failed!\n");
240 return -1;
241 }
242 printf("input mb info: virt = %p, phys = %#lx, size: %d\n", input_virt, input_phys, input_size);
243
244 // Allocate outputs memory in outside
245 MB_BLK output_mbs[io_num.n_output];
246 void* output_virts[io_num.n_output];
247 uint64_t output_physs[io_num.n_output];
248 for (uint32_t i = 0; i < io_num.n_output; ++i) {
249 // default output type is depend on model, this require float32 to compute top5
250 output_attrs[i].type = RKNN_TENSOR_FLOAT32;
251 int output_size = output_attrs[i].n_elems * sizeof(float);
252 output_attrs[i].size = output_size;
253 ret = RK_MPI_MMZ_Alloc(&output_mbs[i], output_size, mb_flags);
254 if (ret < 0) {
255 printf("RK_MPI_MMZ_Alloc failed, ret: %d\n", ret);
256 return ret;
257 }
258 output_virts[i] = RK_MPI_MMZ_Handle2VirAddr(output_mbs[i]);
259 if (output_virts[i] == NULL) {
260 printf("RK_MPI_MMZ_Handle2VirAddr failed!\n");
261 return -1;
262 }
263 output_physs[i] = RK_MPI_MMZ_Handle2PhysAddr(output_mbs[i]);
264 if (output_physs[i] == 0) {
265 printf("RK_MPI_MMZ_Handle2PhysAddr failed!\n");
266 return -1;
267 }
268 printf("output%d mb info: virt = %p, phys = %#lx, size = %d\n", i, output_virts[i], output_physs[i], output_size);
269 }
270
271 // Create input tensor memory
272 rknn_tensor_mem* input_mems[1];
273 // default input type is int8 (normalize and quantize need compute in outside)
274 // if set uint8, will fuse normalize and quantize to npu
275 input_attrs[0].type = input_type;
276 // default fmt is NHWC, npu only support NHWC in zero copy mode
277 input_attrs[0].fmt = input_layout;
278
279 input_mems[0] = rknn_create_mem_from_phys(ctx, input_phys, input_virt, input_attrs[0].size_with_stride);
280
281 // Copy input data to input tensor memory
282 int width = input_attrs[0].dims[2];
283 int stride = input_attrs[0].w_stride;
284 if (width == stride) {
285 memcpy(input_mems[0]->virt_addr, input_data, width * input_attrs[0].dims[1] * input_attrs[0].dims[3]);
286 } else {
287 int height = input_attrs[0].dims[1];
288 int channel = input_attrs[0].dims[3];
289 // copy from src to dst with stride
290 uint8_t* src_ptr = input_data;
291 uint8_t* dst_ptr = (uint8_t*)input_mems[0]->virt_addr;
292 // width-channel elements
293 int src_wc_elems = width * channel;
294 int dst_wc_elems = stride * channel;
295 for (int h = 0; h < height; ++h) {
296 memcpy(dst_ptr, src_ptr, src_wc_elems);
297 src_ptr += src_wc_elems;
298 dst_ptr += dst_wc_elems;
299 }
300 }
301
302 // Create output tensor memory
303 rknn_tensor_mem* output_mems[io_num.n_output];
304 for (uint32_t i = 0; i < io_num.n_output; ++i) {
305 output_mems[i] = rknn_create_mem_from_phys(ctx, output_physs[i], output_virts[i], output_attrs[i].size);
306 }
307
308 // Set input tensor memory
309 ret = rknn_set_io_mem(ctx, input_mems[0], &input_attrs[0]);
310 if (ret < 0) {
311 printf("rknn_set_io_mem fail! ret=%d\n", ret);
312 return -1;
313 }
314
315 // Set output tensor memory
316 for (uint32_t i = 0; i < io_num.n_output; ++i) {
317 // set output memory and attribute
318 ret = rknn_set_io_mem(ctx, output_mems[i], &output_attrs[i]);
319 if (ret < 0) {
320 printf("rknn_set_io_mem fail! ret=%d\n", ret);
321 return -1;
322 }
323 }
324
325 // Run
326 printf("Begin perf ...\n");
327 for (int i = 0; i < loop_count; ++i) {
328 int64_t start_us = getCurrentTimeUs();
329 ret = rknn_run(ctx, NULL);
330 int64_t elapse_us = getCurrentTimeUs() - start_us;
331 if (ret < 0) {
332 printf("rknn run error %d\n", ret);
333 return -1;
334 }
335 printf("%4d: Elapse Time = %.2fms, FPS = %.2f\n", i, elapse_us / 1000.f, 1000.f * 1000.f / elapse_us);
336 }
337
338 // Get top 5
339 uint32_t topNum = 5;
340 for (uint32_t i = 0; i < io_num.n_output; i++) {
341 uint32_t MaxClass[topNum];
342 float fMaxProb[topNum];
343 float* buffer = (float*)output_mems[i]->virt_addr;
344 uint32_t sz = output_attrs[i].n_elems;
345 int top_count = sz > topNum ? topNum : sz;
346
347 rknn_GetTopN(buffer, fMaxProb, MaxClass, sz, topNum);
348
349 printf("---- Top%d ----\n", top_count);
350 for (int j = 0; j < top_count; j++) {
351 printf("%8.6f - %d\n", fMaxProb[j], MaxClass[j]);
352 }
353 }
354
355 // free mb blk memory
356 RK_MPI_MMZ_Free(input_mb);
357 for (uint32_t i = 0; i < io_num.n_output; ++i) {
358 RK_MPI_MMZ_Free(output_mbs[i]);
359 }
360
361 // Destroy rknn memory
362 rknn_destroy_mem(ctx, input_mems[0]);
363 for (uint32_t i = 0; i < io_num.n_output; ++i) {
364 rknn_destroy_mem(ctx, output_mems[i]);
365 }
366
367 // destroy
368 rknn_destroy(ctx);
369
370 free(input_data);
371
372 return 0;
373 }
374