1import numpy as np 2import cv2 3from rknn.api import RKNN 4import torchvision.models as models 5import torch 6import os 7 8 9def show_outputs(output): 10 output_sorted = sorted(output, reverse=True) 11 top5_str = '\n-----TOP 5-----\n' 12 for i in range(5): 13 value = output_sorted[i] 14 index = np.where(output == value) 15 for j in range(len(index)): 16 if (i + j) >= 5: 17 break 18 if value > 0: 19 topi = '{}: {}\n'.format(index[j], value) 20 else: 21 topi = '-1: 0.0\n' 22 top5_str += topi 23 print(top5_str) 24 25 26def show_perfs(perfs): 27 perfs = 'perfs: {}\n'.format(perfs) 28 print(perfs) 29 30 31def softmax(x): 32 return np.exp(x)/sum(np.exp(x)) 33 34def torch_version(): 35 import torch 36 torch_ver = torch.__version__.split('.') 37 torch_ver[2] = torch_ver[2].split('+')[0] 38 return [int(v) for v in torch_ver] 39 40if __name__ == '__main__': 41 42 if torch_version() < [1, 9, 0]: 43 import torch 44 print("Your torch version is '{}', in order to better support the Quantization Aware Training (QAT) model,\n" 45 "Please update the torch version to '1.9.0' or higher!".format(torch.__version__)) 46 exit(0) 47 48 model = './resnet18_i8.pt' 49 50 input_size_list = [[1, 3, 224, 224]] 51 52 # Create RKNN object 53 rknn = RKNN(verbose=True) 54 55 # Pre-process config 56 print('--> Config model') 57 rknn.config(mean_values=[123.675, 116.28, 103.53], std_values=[58.395, 58.395, 58.395]) 58 print('done') 59 60 # Load model 61 print('--> Loading model') 62 ret = rknn.load_pytorch(model=model, input_size_list=input_size_list) 63 if ret != 0: 64 print('Load model failed!') 65 exit(ret) 66 print('done') 67 68 # Build model 69 print('--> Building model') 70 ret = rknn.build(do_quantization=False) 71 if ret != 0: 72 print('Build model failed!') 73 exit(ret) 74 print('done') 75 76 # Export rknn model 77 print('--> Export rknn model') 78 ret = rknn.export_rknn('./resnet_18.rknn') 79 if ret != 0: 80 print('Export rknn model failed!') 81 exit(ret) 82 print('done') 83 84 # Set inputs 85 img = cv2.imread('./space_shuttle_224.jpg') 86 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 87 88 # Init runtime environment 89 print('--> Init runtime environment') 90 ret = rknn.init_runtime() 91 if ret != 0: 92 print('Init runtime environment failed!') 93 exit(ret) 94 print('done') 95 96 # Inference 97 print('--> Running model') 98 outputs = rknn.inference(inputs=[img]) 99 np.save('./pytorch_resnet18_qat_0.npy', outputs[0]) 100 show_outputs(softmax(np.array(outputs[0][0]))) 101 print('done') 102 103 rknn.release() 104