61 lines
2.7 KiB
Python
61 lines
2.7 KiB
Python
import numpy as np
|
|
import cv2
|
|
from rknn.api import RKNN
|
|
|
|
if __name__ == '__main__':
|
|
|
|
# Create RKNN object
|
|
rknn = RKNN(verbose=True)
|
|
|
|
# Pre-process config
|
|
print('--> Config model')
|
|
rknn.config(mean_values=[127.5, 127.5, 127.5], std_values=[127.5, 127.5, 127.5], target_platform='rk3566')
|
|
print('done')
|
|
|
|
# Load model (from https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf1_detection_zoo.md ssd_mobilenet_v2_coco)
|
|
print('--> Loading model')
|
|
ret = rknn.load_tensorflow(tf_pb='./ssd_mobilenet_v2.pb',
|
|
inputs=['FeatureExtractor/MobilenetV2/MobilenetV2/input'],
|
|
outputs=['concat_1', 'concat'],
|
|
input_size_list=[[1,300,300,3]])
|
|
if ret != 0:
|
|
print('Load model failed!')
|
|
exit(ret)
|
|
print('done')
|
|
|
|
# Build model
|
|
print('--> hybrid_quantization_step1')
|
|
ret = rknn.hybrid_quantization_step1(dataset='./dataset.txt', proposal=False)
|
|
if ret != 0:
|
|
print('hybrid_quantization_step1 failed!')
|
|
exit(ret)
|
|
print('done')
|
|
|
|
# Tips
|
|
print("======================================== Tips ==========================================================")
|
|
print("After the execution of 'rknn.hybrid_quantization_step1' is completed, a configuration file")
|
|
print("'ssd_mobilenet_v2.quantization.cfg' will be generated. This configuration file is used to set the layer")
|
|
print("that needs to be hybrid quantized. ")
|
|
print("")
|
|
print("Users can set the 'output name' and 'data type' of the layer that needs to be configured into")
|
|
print("'custom_quantize_layers', as follows:")
|
|
print("")
|
|
print(" custom_quantize_layers:")
|
|
print(" FeatureExtractor/MobilenetV2/expanded_conv/depthwise/BatchNorm/batchnorm/add_1:0: float16")
|
|
print(" FeatureExtractor/MobilenetV2/expanded_conv/depthwise/Relu6:0: float16")
|
|
print("")
|
|
print("Currently supported data types are: 'int8', 'int16', 'float16'.")
|
|
print("")
|
|
print("In addition, the configuration file will automatically fill in some layers that may cause a decrease in")
|
|
print("accuracy (may be inaccurate). Users can consider these layers together with the results of accuracy")
|
|
print("analysis to obtain the layers that actually cause a decrease in accuracy and set them into")
|
|
print("'custom_quantize_layers'.")
|
|
print("")
|
|
print("After modifying 'ssd_mobilenet_v2.quantization.cfg', you can continue to execute 'step2.py' to complete")
|
|
print("the model conversion.")
|
|
print("========================================================================================================")
|
|
print("")
|
|
|
|
rknn.release()
|
|
|