-
Notifications
You must be signed in to change notification settings - Fork 10
/
io.go
38 lines (30 loc) · 862 Bytes
/
io.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
package rknnlite
/*
#include "rknn_api.h"
#include <stdlib.h>
*/
import "C"
import (
"fmt"
"unsafe"
)
// QueryModelIONumber queries the number of Input and Output tensors of the model
func (r *Runtime) QueryModelIONumber() (ioNum IONumber, err error) {
// prepare the structure to receive the Input/Output number
var cIONum C.rknn_input_output_num
// call the C function
ret := C.rknn_query(r.ctx, C.RKNN_QUERY_IN_OUT_NUM, unsafe.Pointer(&cIONum), C.uint(C.sizeof_rknn_input_output_num))
if ret != C.RKNN_SUCC {
return IONumber{}, fmt.Errorf("rknn_query failed with return code %d", int(ret))
}
ioNum = IONumber{
NumberInput: uint32(cIONum.n_input),
NumberOutput: uint32(cIONum.n_output),
}
return ioNum, nil
}
// IONumber represents the C.rknn_input_output_num struct
type IONumber struct {
NumberInput uint32
NumberOutput uint32
}