forked from Github/frigate
rockchip: update dependencies and add script for model conversion (#15699)
* rockchip: update dependencies and add script for model conversion * rockchip: update docs --------- Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>
This commit is contained in:
committed by
Nicolas Mowen
parent
39ecffbf92
commit
e8d27d1f91
82
docker/rockchip/conv2rknn.py
Normal file
82
docker/rockchip/conv2rknn.py
Normal file
@@ -0,0 +1,82 @@
|
||||
import os
|
||||
|
||||
import rknn
|
||||
import yaml
|
||||
from rknn.api import RKNN
|
||||
|
||||
try:
|
||||
with open(rknn.__path__[0] + "/VERSION") as file:
|
||||
tk_version = file.read().strip()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
try:
|
||||
with open("/config/conv2rknn.yaml", "r") as config_file:
|
||||
configuration = yaml.safe_load(config_file)
|
||||
except FileNotFoundError:
|
||||
raise Exception("Please place a config.yaml file in /config/conv2rknn.yaml")
|
||||
|
||||
if configuration["config"] != None:
|
||||
rknn_config = configuration["config"]
|
||||
else:
|
||||
rknn_config = {}
|
||||
|
||||
if not os.path.isdir("/config/model_cache/rknn_cache/onnx"):
|
||||
raise Exception(
|
||||
"Place the onnx models you want to convert to rknn format in /config/model_cache/rknn_cache/onnx"
|
||||
)
|
||||
|
||||
if "soc" not in configuration:
|
||||
try:
|
||||
with open("/proc/device-tree/compatible") as file:
|
||||
soc = file.read().split(",")[-1].strip("\x00")
|
||||
except FileNotFoundError:
|
||||
raise Exception("Make sure to run docker in privileged mode.")
|
||||
|
||||
configuration["soc"] = [
|
||||
soc,
|
||||
]
|
||||
|
||||
if "quantization" not in configuration:
|
||||
configuration["quantization"] = False
|
||||
|
||||
if "output_name" not in configuration:
|
||||
configuration["output_name"] = "{{input_basename}}"
|
||||
|
||||
for input_filename in os.listdir("/config/model_cache/rknn_cache/onnx"):
|
||||
for soc in configuration["soc"]:
|
||||
quant = "i8" if configuration["quantization"] else "fp16"
|
||||
|
||||
input_path = "/config/model_cache/rknn_cache/onnx/" + input_filename
|
||||
input_basename = input_filename[: input_filename.rfind(".")]
|
||||
|
||||
output_filename = (
|
||||
configuration["output_name"].format(
|
||||
quant=quant,
|
||||
input_basename=input_basename,
|
||||
soc=soc,
|
||||
tk_version=tk_version,
|
||||
)
|
||||
+ ".rknn"
|
||||
)
|
||||
output_path = "/config/model_cache/rknn_cache/" + output_filename
|
||||
|
||||
rknn_config["target_platform"] = soc
|
||||
|
||||
rknn = RKNN(verbose=True)
|
||||
rknn.config(**rknn_config)
|
||||
|
||||
if rknn.load_onnx(model=input_path) != 0:
|
||||
raise Exception("Error loading model.")
|
||||
|
||||
if (
|
||||
rknn.build(
|
||||
do_quantization=configuration["quantization"],
|
||||
dataset="/COCO/coco_subset_20.txt",
|
||||
)
|
||||
!= 0
|
||||
):
|
||||
raise Exception("Error building model.")
|
||||
|
||||
if rknn.export_rknn(output_path) != 0:
|
||||
raise Exception("Error exporting rknn model.")
|
||||
Reference in New Issue
Block a user