export_pytorch_model() model = './resnet18.pt' input_size_list = [[3,224,224]] # Create RKNN object rknn = RKNN() # pre-process config print('--> config model') rknn.config(channel_mean_value='123.675 116.28 103.53 58.395', reorder_channel='0 1 2') print('done') # Load pytorch model print('--> Loading model') ret = rknn.load_pytorch(model=model, input_size_list=input_size_list) if ret != 0: print('Load pytorch model failed!') exit(ret) print('done') # Build model print('--> Building model') ret = rknn.build(do_quantization=True, dataset='./dataset.txt') if ret != 0: print('Build pytorch failed!') exit(ret) print('done') # Export rknn model print('--> Export RKNN model')
if __name__ == '__main__': # Create RKNN object rknn = RKNN() # model config print('--> Config model') rknn.config(mean_values=[[123.675, 116.28, 103.53]], std_values=[[58.395, 58.395, 58.395]], reorder_channel='0 1 2') print('done') # Load Pytorch model print('--> Loading model') ret = rknn.load_pytorch(model='./mnasnet0_5.pt', input_size_list=[[3, 224, 224]]) if ret != 0: print('Load model failed!') exit(ret) print('done') # Hybrid quantization step1 print('--> hybrid_quantization_step1') ret = rknn.hybrid_quantization_step1(dataset='./dataset.txt') if ret != 0: print('hybrid_quantization_step1 failed!') exit(ret) print('done') # Tips print('Please modify mnasnet0_5.quantization.cfg!')