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!')