深入探索 RKNN 模型轉換之旅

在人工智能蓬勃發展的當下，邊緣計算領域的應用愈發廣泛。瑞芯微的 RKNN 技術在這一領域大放異彩，它能讓深度學習模型在其芯片平臺上高效運行。而在整個應用流程中，模型轉換是極為關鍵的一環，今天就讓我們一同深入這個神奇的 RKNN 模型轉換世界。?

一、從常見模型邁向 ONNX?

在深度學習的江湖中，存在著眾多模型框架，如 PyTorch、TensorFlow 等。這些模型就如同身懷絕技的大俠，各有千秋。但要讓它們能在瑞芯微的平臺上施展拳腳，往往需要先將其轉換為 ONNX 格式作為中間橋梁。?

（一）PyTorch 模型轉 ONNX?

假設我們有一個訓練好的 PyTorch 模型，要將其轉換為 ONNX 格式，首先要確保我們的環境中安裝模型運行所依賴庫以及onnx庫。我們以yolov5官方提供的export.py來講解

我們想來看export.py中的run函數（我們只關注onnx轉化部分）

def run(data=ROOT / 'data/coco128.yaml',  # 'dataset.yaml path'weights=ROOT / 'yolov5s.pt',  # weights pathimgsz=(640, 640),  # image (height, width)batch_size=1,  # batch sizedevice='cpu',  # cuda device, i.e. 0 or 0,1,2,3 or cpuinclude=('torchscript', 'onnx', 'coreml'),  # include formatshalf=False,  # FP16 half-precision exportinplace=False,  # set YOLOv5 Detect() inplace=Truetrain=False,  # model.train() modeoptimize=False,  # TorchScript: optimize for mobileint8=False,  # CoreML/TF INT8 quantizationdynamic=False,  # ONNX/TF: dynamic axessimplify=False,  # ONNX: simplify modelopset=12,  # ONNX: opset versionverbose=False,  # TensorRT: verbose logworkspace=4,  # TensorRT: workspace size (GB)topk_per_class=100,  # TF.js NMS: topk per class to keeptopk_all=100,  # TF.js NMS: topk for all classes to keepiou_thres=0.45,  # TF.js NMS: IoU thresholdconf_thres=0.25  # TF.js NMS: confidence threshold):t = time.time()include = [x.lower() for x in include]tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'tfjs'))  # TensorFlow exportsimgsz *= 2 if len(imgsz) == 1 else 1  # expandfile = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)# Load PyTorch modeldevice = select_device(device)assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'model = attempt_load(weights, map_location=device, inplace=True, fuse=True)  # load FP32 modelnc, names = model.nc, model.names  # number of classes, class names# Inputgs = int(max(model.stride))  # grid size (max stride)imgsz = [check_img_size(x, gs) for x in imgsz]  # verify img_size are gs-multiplesim = torch.zeros(batch_size, 3, *imgsz).to(device)  # image size(1,3,320,192) BCHW iDetection# Update modelif half:im, model = im.half(), model.half()  # to FP16model.train() if train else model.eval()  # training mode = no Detect() layer grid constructionfor k, m in model.named_modules():if isinstance(m, Conv):  # assign export-friendly activationsif isinstance(m.act, nn.SiLU):m.act = SiLU()elif isinstance(m, Detect):m.inplace = inplacem.onnx_dynamic = dynamic# m.forward = m.forward_export  # assign forward (optional)for _ in range(2):y = model(im)  # dry runsLOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")# Exportsif 'torchscript' in include:export_torchscript(model, im, file, optimize)if 'onnx' in include:export_onnx(model, im, file, opset, train, dynamic, simplify)if 'engine' in include:export_engine(model, im, file, train, half, simplify, workspace, verbose)if 'coreml' in include:export_coreml(model, im, file)# TensorFlow Exportsif any(tf_exports):pb, tflite, tfjs = tf_exports[1:]assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'model = export_saved_model(model, im, file, dynamic, tf_nms=tfjs, agnostic_nms=tfjs,topk_per_class=topk_per_class, topk_all=topk_all, conf_thres=conf_thres,iou_thres=iou_thres)  # keras modelif pb or tfjs:  # pb prerequisite to tfjsexport_pb(model, im, file)if tflite:export_tflite(model, im, file, int8=int8, data=data, ncalib=100)if tfjs:export_tfjs(model, im, file)# FinishLOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'f"\nResults saved to {colorstr('bold', file.parent.resolve())}"f'\nVisualize with https://netron.app')True)

參數解析

在run函數中，有許多參數用于控制模型導出的過程，下面我們來詳細了解一下這些參數的作用：

data：數據集配置文件的路徑，一般包含數據集的相關信息，在導出模型時可能會用于一些數據相關的配置，但在單純的模型格式轉換過程中，它并非關鍵參數。
weights：要導出的 PyTorch 模型權重文件路徑，例如yolov5s.pt，這是我們轉換的源模型。在代碼中，通過Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)處理權重文件路徑，如果路徑是 HTTP 或 HTTPS 開頭，會先將其轉換為本地文件路徑。
imgsz：指定輸入圖像的尺寸，格式為元組(height, width)，例如(640, 640)，模型在推理時會將輸入圖像調整為該尺寸。代碼中gs = int(max(model.stride))獲取模型的最大步長，imgsz = [check_img_size(x, gs) for x in imgsz]確保輸入圖像尺寸是步長的倍數，保證模型結構與輸入尺寸匹配。

gs = int(max(model.stride))  # grid size (max stride)
imgsz = [check_img_size(x, gs) for x in imgsz]  # verify img_size are gs-multiples
im = torch.zeros(batch_size, 3, *imgsz).to(device)  # image size(1,3,320,192) BCHW iDetection這個輸入張量用于模擬實際推理時的輸入數據，在模型導出過程中，PyTorch 會根據這個輸入跟蹤模型的計算圖，從而生成對應的 ONNX 模型。

im = torch.zeros(batch_size, 3, *imgsz).to(device)這個輸入張量用于模擬實際推理時的輸入數據，在模型導出過程中，PyTorch 會根據這個輸入跟蹤模型的計算圖，從而生成對應的 ONNX 模型。?

batch_size：輸入圖像的批量大小，默認為1，即一次處理一張圖像。在實際部署中，可以根據硬件資源和需求調整該值。后續通過im = torch.zeros(batch_size, 3, *imgsz).to(device)創建對應批量大小的輸入張量。
device：指定模型運行和導出的設備，可以是cpu或者cuda設備編號（如0、0,1,2,3等）。select_device(device)函數負責選擇合適的設備，并通過assert not (device.type == 'cpu' and half)限制半精度導出只能在 GPU 設備上進行，具體代碼如下。

device = select_device(device)
assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
model = attempt_load(weights, map_location=device, inplace=True, fuse=True)  
# load FP32 model inplace=True表示在原模型上進行操作以節省內存，fuse=True會對模型中的一些層進行融合優化，提高推理速度。

include：一個元組，用于指定要導出的模型格式，其中'onnx'就是我們關注的將模型轉換為 ONNX 格式的標識。代碼通過include = [x.lower() for x in include]將格式名稱統一轉換為小寫，方便后續判斷。
half：是否以 FP16 半精度格式導出模型，需要注意的是，半精度導出僅在 GPU 設備上支持。當half=True時，通過im, model = im.half(), model.half()將輸入數據和模型轉換為 FP16 格式。
dynamic：對于 ONNX 格式，該參數表示是否使用動態軸，即輸入的尺寸等維度可以在推理時動態變化，增加了模型使用的靈活性。在后續export_onnx函數中會根據該參數設置動態軸。

if half:im, model = im.half(), model.half()  # to FP16
model.train() if train else model.eval()  # training mode = no Detect() layer grid construction
for k, m in model.named_modules():if isinstance(m, Conv):  # assign export-friendly activationsif isinstance(m.act, nn.SiLU):m.act = SiLU()elif isinstance(m, Detect):m.inplace = inplacem.onnx_dynamic = dynamic# m.forward = m.forward_export  # assign forward (optional)for _ in range(2):y = model(im)  # dry runs

當half=True時，通過im.half()和model.half()將輸入數據im和模型model都轉換為 FP16 半精度。這樣做可以減少模型的內存占用和推理計算量，但僅在 GPU 設備上支持。根據train參數決定將模型設置為訓練模式（model.train()）還是評估模式（model.eval()），在導出模型用于推理時，通常設置為評估模式，此時 YOLOv5 模型中的檢測層不會進行訓練時的網格構建等操作。

通過遍歷模型的各個模塊for k, m in model.named_modules()，對卷積層（isinstance(m, Conv)）進行處理。如果卷積層的激活函數是nn.SiLU，將其替換為SiLU()，這是因為nn.SiLU在導出為 ONNX 格式時可能存在兼容性問題，替換為自定義的SiLU激活函數可以提高導出成功率。對于檢測層（isinstance(m, Detect)），設置m.inplace = inplace和m.onnx_dynamic = dynamic，前者控制檢測層是否在原地進行操作，后者設置 ONNX 導出時是否使用動態軸。

最后進行兩次model(im)的前向傳播（干運行，dry runs），目的是讓模型的一些參數和緩存進行初始化，確保在正式導出模型時計算圖的穩定和正確，避免因參數未初始化導致的導出錯誤。

simplify：是否對導出的 ONNX 模型進行簡化，簡化后的模型可能在推理時具有更好的性能和效率。在export_onnx函數中會根據此參數決定是否調用onnx-simplifier庫進行模型簡化。
opset：指定 ONNX 的算子集版本，不同版本的算子集對模型的支持和兼容性有所不同，export.py中默認使用12版本。該參數在export_onnx函數中傳遞給torch.onnx.export，影響模型轉換過程中算子的映射和表示（這里需要注意我們下文在轉化為rknn模型是使用的環境是rknn-toolkit2的1.5.2版本，該版本最高支持opset版本為12）。

if 'onnx' in include:export_onnx(model, im, file, opset, train, dynamic, simplify)

最后模型調用?export_onnx函數進行轉換，我們接下倆看這個函數具體實現

def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):# YOLOv5 ONNX exporttry:check_requirements(('onnx',))import onnxLOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')f = file.with_suffix('.onnx')torch.onnx.export(model, im, f, verbose=False, opset_version=opset,training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,do_constant_folding=not train,input_names=['images'],output_names=['output'],dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)} if dynamic else None)# Checksmodel_onnx = onnx.load(f)  # load onnx modelonnx.checker.check_model(model_onnx)  # check onnx model# LOGGER.info(onnx.helper.printable_graph(model_onnx.graph))  # print# Simplifyif simplify:try:check_requirements(('onnx-simplifier',))import onnxsimLOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')model_onnx, check = onnxsim.simplify(model_onnx,dynamic_input_shape=dynamic,input_shapes={'images': list(im.shape)} if dynamic else None)assert check, 'assert check failed'onnx.save(model_onnx, f)except Exception as e:LOGGER.info(f'{prefix} simplifier failure: {e}')LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')LOGGER.info(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'")except Exception as e:LOGGER.info(f'{prefix} export failure: {e}')

?這一部分代碼中核心是torch.onnx.export這個api的使用，我們先來展開它的參數

torch.onnx.export(model, im, f, verbose=False, opset_version=opset,training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,do_constant_folding=not train,input_names=['images'],output_names=['output'],dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)} if dynamic else None)

model：要導出的 PyTorch 模型，即前面run函數中加載并處理后的 YOLOv5 模型。它是整個導出過程的核心對象，后續的操作都是圍繞這個模型的計算圖進行轉換。
im：模型的輸入張量，在run函數中通過torch.zeros(batch_size, 3, *imgsz).to(device)創建。這個輸入張量用于追蹤模型的計算過程，PyTorch 會根據它在模型中的流動路徑構建計算圖，并將其轉換為 ONNX 格式的計算圖。簡單來說，它就像是給模型一個 “示例輸入”，告訴系統模型在實際運行時輸入數據的形狀和類型。
f：指定導出的 ONNX 模型保存路徑，是一個Path對象，由file.with_suffix('.onnx')生成。例如，如果原始權重文件是yolov5s.pt，那么f可能是yolov5s.onnx，模型最終會以這個文件名保存在相應目錄下。
verbose=False：該參數控制導出過程中是否輸出詳細的日志信息。當設置為False時，導出過程較為安靜，不會打印大量中間信息；如果設置為True，則會輸出更多關于模型轉換過程的細節，通常在調試導出問題時會將其設置為True。
opset_version=opset：opset即算子集版本，它決定了在轉換過程中使用哪些 ONNX 算子來表示 PyTorch 模型中的操作。不同的算子集版本對算子的支持和實現方式有所不同，這里使用run函數傳入的opset參數（默認值為 12），以確保模型中的操作能正確映射到 ONNX 算子上。如果模型中使用的某些操作在選定的算子集版本中不支持，就會導致導出失敗（同樣下文轉化rknn模型我們使用的環境最高支持12版本，注意不要使用13或更高版本）。
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL：根據train參數的值來設置模型導出時的模式。如果train=True，則將模型設置為訓練模式導出，此時模型中的一些操作（如 BatchNorm 層的計算方式）會按照訓練模式進行轉換；如果train=False（通常用于導出推理模型），則設置為評估模式，模型以推理時的行為進行轉換，比如 BatchNorm 層會使用固定的均值和方差。
do_constant_folding=not train：常量折疊是一種優化技術，用于在導出過程中簡化模型。當train=False（即導出推理模型）時，do_constant_folding=True，會將模型中一些固定的常量計算提前完成，減少推理時的計算量。例如，如果模型中有x = 2 + 3這樣的固定計算，常量折疊會直接將其替換為x = 5。在訓練模式下，由于模型參數可能會變化，一般不進行常量折疊。
input_names=['images']：為導出的 ONNX 模型的輸入指定名稱。這里將輸入命名為'images'，在后續使用 ONNX 模型進行推理時，可以根據這個名稱來提供正確的輸入數據。這個名稱主要用于標識輸入的用途，方便與外部接口進行數據交互。
output_names=['output']：類似地，為導出的 ONNX 模型的輸出指定名稱為'output'。在推理時，通過這個名稱可以從模型輸出中獲取推理結果。它定義了模型輸出數據在 ONNX 格式中的標識。
dynamic_axes：這是一個關鍵參數，用于設置動態軸。當dynamic=True時，會為輸入和輸出張量指定哪些維度是動態的。例如，對于輸入'images'，{0: 'batch', 2: 'height', 3: 'width'}表示第 0 維（批量大小）、第 2 維（圖像高度）和第 3 維（圖像寬度）是可以變化的，在推理時可以傳入不同批量大小或不同尺寸的圖像；對于輸出'output'，{0: 'batch', 1: 'anchors'}表示輸出張量的第 0 維（批量大小）和第 1 維（錨框相關維度）是動態的。如果dynamic=False，則不使用動態軸，輸入和輸出張量的形狀在導出時就固定下來，推理時只能使用固定形狀的輸入數據?

model_onnx = onnx.load(f)  # load onnx model
onnx.checker.check_model(model_onnx)  # check onnx model

onnx.load(f)：使用onnx庫的load函數加載剛剛導出的 ONNX 模型文件。這一步將磁盤上的.onnx文件讀取到內存中，生成一個model_onnx對象，后續可以對這個對象進行各種操作和檢查。
onnx.checker.check_model(model_onnx)：調用onnx.checker模塊的check_model函數對加載的 ONNX 模型進行合法性檢查。它會檢查模型的結構是否符合 ONNX 規范，比如節點連接是否正確、數據類型是否匹配、算子使用是否合規等。如果模型存在問題，會拋出相應的錯誤信息，提示用戶模型導出過程中可能出現了問題，需要進一步排查。

if simplify:try:check_requirements(('onnx-simplifier',))import onnxsimLOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')model_onnx, check = onnxsim.simplify(model_onnx,dynamic_input_shape=dynamic,input_shapes={'images': list(im.shape)} if dynamic else None)assert check, 'assert check failed'onnx.save(model_onnx, f)except Exception as e:LOGGER.info(f'{prefix} simplifier failure: {e}')

?當run函數傳入的simplify參數為True時，會執行模型簡化操作：

環境檢查：check_requirements(('onnx-simplifier',))檢查當前環境是否安裝了onnx-simplifier庫，如果沒有安裝則提示安裝依賴。onnx-simplifier庫專門用于對 ONNX 模型進行優化和簡化。
導入庫與日志記錄：import onnxsim導入onnx-simplifier庫，LOGGER.info記錄日志，提示開始使用onnx-simplifier對模型進行簡化，并打印出當前庫的版本號。
模型簡化：onnxsim.simplify函數用于執行簡化操作，傳入加載的model_onnx對象、dynamic_input_shape=dynamic（根據dynamic參數決定是否按照動態軸進行簡化）以及輸入形狀信息（如果是動態軸，傳入input_shapes={'images': list(im.shape)}，指定輸入張量的形狀）。函數返回簡化后的模型model_onnx以及一個檢查結果check，check用于判斷簡化過程是否成功。
保存簡化后模型：通過assert check, 'assert check failed'確保簡化過程成功，如果失敗則拋出斷言錯誤。最后使用onnx.save(model_onnx, f)將簡化后的模型覆蓋保存到原來的文件路徑f下，完成模型簡化和保存的過程。如果簡化過程中出現異常，會捕獲異常并通過LOGGER.info記錄錯誤信息，提示用戶模型簡化失敗。

關于pytorch計算圖和onnx轉化的更多細節可以參考Pytorch 轉ONNX詳解

（二）TensorFlow模型轉 ONNX?

對于 TensorFlow 模型，轉換過程也有其獨特的步驟，并且TensorFlow1.x和TensorFlow2.x步驟有所區別，由于我沒有使用過TensorFlow2進行轉化，所以我只討論TensorFlow1.x的轉化。

關于TensorFlow2.x的轉化可以參考將tensorflow 1.x & 2.x轉化成onnx文件（以arcface-tf2人臉識別模型為例）（我下文TensorFlow1的轉化同樣參考這篇文章做的）

TensorFlow 1.x 與 ONNX 之間的轉換本質是計算圖的映射過程：

計算圖提取：從 TensorFlow 的會話 (Session) 中提取計算圖結構和權重參數
算子映射：將 TensorFlow 特定算子轉換為 ONNX 標準算子
元數據轉換：保留輸入輸出名稱、數據類型等元信息
模型序列化：生成符合 ONNX 規范的二進制文件

這種轉換使得模型可以在支持 ONNX 的推理引擎 (如 TensorRT、OpenVINO、ONNX Runtime) 上高效運行。

首先要安裝tensorflow、onnx和tf2onnx庫。

def convert_to_onnx(checkpoint_path, output_onnx_path, img_size, action_dim, scope_name, is_relu):if not os.path.exists(checkpoint_path):os.makedirs(checkpoint_path)tf.reset_default_graph()model = network(img_size=img_size, myScope=scope_name)init = tf.global_variables_initializer()saver = tf.train.Saver(max_to_keep=None)trainables = tf.trainable_variables()

環境準備：
- tf.reset_default_graph()：清除默認計算圖，避免干擾
- os.makedirs(checkpoint_path)：確保檢查點路徑存在
模型初始化：
- 創建網絡實例
- 定義變量初始化器和模型保存器

with tf.Session() as sess:sess.run(init)ckpt = tf.train.get_checkpoint_state(checkpoint_path)saver.restore(sess, ckpt.model_checkpoint_path)     input_names = ["Input_1:0","Input_2:0"]output_names = ["Output:0"]

會話管理：
- 創建 TensorFlow 會話并初始化變量
- 從檢查點恢復模型權重
輸入輸出命名：可以通過model.Input1.name打印輸入輸出名

frozen_graph = tf2onnx.tf_loader.freeze_session(sess,input_names=input_names,output_names=output_names
)

凍結計算圖的作用：
- 將變量值嵌入到計算圖中，形成靜態圖
- 移除與推理無關的操作 (如訓練節點)
- 生成適合導出的固化模型

onnx_model = tf2onnx.convert.from_graph_def(frozen_graph,input_names=input_names,output_names=output_names,opset=12,targeted_onnx=onnx.defs.onnx_opset_version(),large_model=False,output_path=output_onnx_path,convert_to="float16"  # 指定半精度轉換
)with open(output_onnx_path, "wb") as f:f.write(onnx_model[0].SerializeToString())
print(f"模型已保存到 {output_onnx_path}")

關鍵參數：
- frozen_graph：凍結后的計算圖
- input_names/output_names：輸入輸出節點名稱
- opset=12：指定 ONNX 算子集版本
轉換流程：
1. 解析 TensorFlow 計算圖
2. 將 TensorFlow 算子映射到 ONNX 算子
3. 生成 ONNX 模型對象
4. 序列化為二進制文件

onnx_model = onnx.load(output_onnx_path)
onnx.checker.check_model(onnx_model)
print("======================================================")
print("success!")

驗證步驟：
- 加載導出的 ONNX 模型
- 檢查模型結構是否符合 ONNX 規范
- 驗證節點連接、數據類型和張量形狀

下面給出完整轉化代碼（網絡細節需要填充，并且這份代碼網絡有多個輸入）

import tensorflow as tf
import tf2onnx
import onnx
class network():def __init__(self, img_size, myScope):pass #這里填充網絡結構def convert_to_onnx(checkpoint_path, output_onnx_path, img_size, action_dim, scope_name,is_relu):if not os.path.exists(checkpoint_path):os.makedirs(checkpoint_path)tf.reset_default_graph()model = network(img_size=img_size， myScope=scope_name)init = tf.global_variables_initializer()saver = tf.train.Saver(max_to_keep=None)trainables = tf.trainable_variables()with tf.Session() as sess:sess.run(init)ckpt = tf.train.get_checkpoint_state(checkpoint_path)saver.restore(sess, ckpt.model_checkpoint_path)     input_names = ["Input_1:0","Input_2:0"]output_names = ["Output:0"]frozen_graph = tf2onnx.tf_loader.freeze_session(sess,input_names=input_names,output_names=output_names)onnx_model = tf2onnx.convert.from_graph_def(frozen_graph,input_names=input_names,output_names=output_names,opset=12,targeted_onnx=onnx.defs.onnx_opset_version(),large_model=False,output_path=output_onnx_path,convert_to="float16"  # 指定半精度轉換)with open(output_onnx_path, "wb") as f:f.write(onnx_model[0].SerializeToString())print(f"模型已保存到 {output_onnx_path}")onnx_model = onnx.load(output_onnx_path)onnx.checker.check_model(onnx_model)print("======================================================")print("success!")if __name__ == "__main__":convert_to_onnx(checkpoint_path="./path",     output_onnx_path="onnx_model.onnx",img_size=640,scope_name="main",  # 與 Qnetwork 初始化時的 myScope 參數一致is_relu=False)

這一步其實可以直接進行模型量化

# 使用 onnxruntime 進行量化
from onnxruntime.quantization import quantize_dynamic, QuantTypequantized_model = quantize_dynamic(output_onnx_path,"quantized_model.onnx",weight_type=QuantType.QInt8
)

但我沒有測試過這種，不太確定。?

二、借助 Docker 掛載 rknn - toolkit2 鏡像

完成模型到 ONNX 格式的轉換后，接下來要進入一個新的階段 —— 利用 rknn - toolkit2 將 ONNX 模型轉換為 RKNN 模型。為了讓這個過程更加順利和便捷，我們可以借助 Docker 來掛載 rknn - toolkit2 鏡像。?

（一）安裝 Docker?

如果你的系統中還沒有安裝 Docker，那么第一步就是安裝它。以 Ubuntu 系統為例，在終端中輸入以下命令：

sudo apt update
wget http://fishros.com/install -O fishros && . fishros

這兩條命令會先更新軟件包列表，然后按照終端中提示的數字選擇安裝 Docker。安裝完成后，可以通過sudo docker run hello - world命令來驗證 Docker 是否安裝成功。如果能看到 Docker 輸出的一些歡迎信息，那就說明安裝大功告成啦。?

（二）下載 rknn - toolkit2 鏡像?

瑞芯微官方提供了 rknn - toolkit2 的 Docker 鏡像，我們可以從瑞芯微官方github上下載。（我使用的是更早版本1.5.2，下文也是在此版本上轉化）

下載完成后，進入rknn-toolkit2-{版本號}.tar.gz文件夾，運行

docker load -i ./rknn-toolkit2-（版本號）-cp（不同版本不同）-docker.tar.gz
以1.5.2為例，命令如下
docker load -i ./rknn-toolkit2-1.5.2-cp36-docker.tar.gz

我們可以通過sudo docker images命令查看本地已有的鏡像列表，確認 rknn - toolkit2 鏡像是否已成功下載。?

（三）掛載并運行鏡像?

當鏡像下載好后，就可以進行掛載和運行了。假設我們有一個存放 ONNX 模型和相關轉換代碼的目錄/home/user/rknn_convert，我們希望將這個目錄掛載到 Docker 容器中，并在容器內運行轉換操作。可以使用進入這個目錄并運行以下命令：

docker run  -it --name rknn_toolkit2 -v $(pwd):/app rknn-toolkit2:1.5.2-cp36

這里的-t -i參數表示以交互式終端的方式運行容器，-v $(pwd):/app表示將本地的當前路徑掛載到容器內的/app目錄，這樣在容器內就可以訪問和操作我們本地的模型和代碼了，最后指定鏡像名稱。

輸入：

sudo docker ps

查看正在運行的docker容器

?在vscode中安裝插件ms-vscode-remote.remote-containers，安裝好后，ctrl+shift+p，彈出窗口輸入attach，下拉框中選擇正在運行的容器docker，并選擇我們之前運行的容器，就可以用vscode在容器內編寫代碼了。

三、ONNX 模型華麗變身 RKNN 模型?

當我們成功進入掛載了 rknn - toolkit2 鏡像的 Docker 容器后，就可以正式開始將 ONNX 模型轉換為 RKNN 模型的神奇之旅了。?

（一）整體流程概覽

將 ONNX 模型轉換為 RKNN 模型，主要分為以下幾個步驟：

環境準備：創建 RKNN 對象，配置預處理參數。
模型加載：將 ONNX 格式的模型文件加載到 RKNN 環境中。
模型構建：這一步包含了關鍵的模型量化操作，將模型轉換為適合 RKNN 推理框架運行的形式。
模型導出：生成最終可在瑞芯微設備上運行的 RKNN 模型文件。
推理驗證：在本地初始化運行環境，輸入測試圖像進行推理，并對結果進行后處理展示。

接下來，我們結合代碼逐步深入解析每個步驟。

（二）代碼實現詳解

1. 環境配置與初始化

# Create RKNN object
rknn = RKNN(verbose=True)# pre-process config
print('--> Config model')
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform='rk3588')
print('done')

首先，通過RKNN(verbose=True)創建一個 RKNN 對象，并設置verbose=True開啟詳細日志輸出，方便我們在后續操作中排查問題。

然后，使用rknn.config()方法進行預處理配置。mean_values和std_values用于對輸入圖像進行歸一化處理。這里mean_values=[[0, 0, 0]]（不填取默認全0）表示將圖像每個通道的像素值減去 0，std_values=[[255, 255, 255]]（不填取默認全1）表示將結果除以 255，相當于把像素值從[0, 255]的范圍映射到[0, 1]。target_platform='rk3588'則指定了模型要部署的目標硬件平臺為瑞芯微 RK3588，不同的平臺可能對模型的優化策略有所不同。

2. 加載 ONNX 模型

# Load ONNX model
print('--> Loading model')
ret = rknn.load_onnx(model=ONNX_MODEL)
if ret != 0:print('Load model failed!')exit(ret)
print('done')

rknn.load_onnx(model=ONNX_MODEL)這行代碼負責將指定路徑下的 ONNX 模型文件（ONNX_MODEL變量定義了文件路徑，這里假設為'yolov5s.onnx'）加載到 RKNN 對象中。如果加載過程返回值不為 0，說明加載失敗，程序會輸出錯誤提示并終止運行，以避免后續無效操作?

3. 模型構建與量化

# Build model
print('--> Building model')
ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
if ret != 0:print('Build model failed!')exit(ret)
print('done')

rknn.build()是整個轉換過程的核心，其中do_quantization=QUANTIZE_ON參數決定是否進行模型量化。QUANTIZE_ON是一個全局布爾變量，在代碼開頭定義為True，即開啟量化。

為什么要進行模型量化？
在嵌入式設備上，計算資源和內存空間相對有限。模型量化通過將模型參數和計算從高精度（如 32 位浮點數）轉換為低精度（如 8 位整數），可以大幅減少模型的內存占用和計算量，從而提升模型在設備上的運行速度和效率。例如，原本用 32 位浮點數表示的權重參數，量化后用 8 位整數表示，數據存儲空間直接縮小為原來的 1/4。

量化過程中的數據集作用
dataset=DATASET參數傳入了一個數據集文件（假設路徑為'./datasets.txt'，至少有一個樣本），這個數據集在量化過程中扮演著重要角色。在量化時，RKNN 需要通過分析數據集中的樣本數據，確定如何將高精度數據映射到低精度數據，以盡可能減少量化帶來的精度損失。datasets.txt文件通常是一個文本文件，每一行記錄一個圖像文件的路徑，RKNN 會讀取這些圖像，按照之前配置的預處理方式進行處理，然后用處理后的數據來校準模型量化的參數，比如確定合適的量化縮放因子和零點偏移等。

如果模型構建（包括量化）過程返回值不為 0，同樣說明操作失敗，程序會輸出錯誤信息并終止。

?4. 導出 RKNN 模型

# Export RKNN model
print('--> Export rknn model')
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:print('Export rknn model failed!')exit(ret)
print('done')

模型構建完成后，使用rknn.export_rknn(RKNN_MODEL)將轉換并量化后的模型導出為 RKNN 格式的文件（RKNN_MODEL變量指定了輸出文件路徑，這里假設為'yolov5s_quant.rknn'）。如果導出失敗，程序會提示錯誤并退出。導出成功后，得到的.rknn文件就可以部署到瑞芯微的目標設備上運行了。

5. 推理驗證

# Init runtime environment
print('--> Init runtime environment')
ret = rknn.init_runtime()
# ret = rknn.init_runtime('rk3566')
if ret != 0:print('Init runtime environment failed!')exit(ret)
print('done')# Set inputs
img = cv2.imread(IMG_PATH)
# img, ratio, (dw, dh) = letterbox(img, new_shape=(IMG_SIZE, IMG_SIZE))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))# Inference
print('--> Running model')
outputs = rknn.inference(inputs=[img])
np.save('./onnx_yolov5_0.npy', outputs[0])
np.save('./onnx_yolov5_1.npy', outputs[1])
np.save('./onnx_yolov5_2.npy', outputs[2])
print('done')# post process
input0_data = outputs[0]
input1_data = outputs[1]
input2_data = outputs[2]input0_data = input0_data.reshape([3, -1]+list(input0_data.shape[-2:]))
input1_data = input1_data.reshape([3, -1]+list(input1_data.shape[-2:]))
input2_data = input2_data.reshape([3, -1]+list(input2_data.shape[-2:]))input_data = list()
input_data.append(np.transpose(input0_data, (2, 3, 0, 1)))
input_data.append(np.transpose(input1_data, (2, 3, 0, 1)))
input_data.append(np.transpose(input2_data, (2, 3, 0, 1)))boxes, classes, scores = yolov5_post_process(input_data)img_1 = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
if boxes is not None:draw(img_1, boxes, scores, classes)cv2.imwrite('result.png', img_1)rknn.release()

初始化運行環境：rknn.init_runtime()用于初始化 RKNN 的運行環境，為模型推理做好準備。如果指定目標平臺，也可以傳入平臺名稱，如rknn.init_runtime('rk3566')?。
設置輸入數據：讀取測試圖像（路徑由IMG_PATH指定），進行顏色空間轉換（從 BGR 轉換為 RGB，因為模型訓練時可能使用的是 RGB 格式）和尺寸縮放（調整為模型輸入要求的IMG_SIZE，這里是640x640）。
執行推理：rknn.inference(inputs=[img])使用加載并轉換好的模型對輸入圖像進行推理，得到輸出結果outputs。這里將輸出結果保存為.npy文件，方便后續分析。
后處理與結果展示：對推理輸出進行一系列后處理操作，包括調整數據形狀、轉換維度順序，然后通過yolov5_post_process函數進行目標檢測結果的解析，得到檢測框的位置、類別和置信度等信息。最后使用draw函數將檢測結果繪制在原始圖像上，并保存為result.png。
資源釋放：rknn.release()釋放 RKNN 占用的資源，結束整個模型轉換與推理流程。

6.量化過程中的關鍵要點與注意事項

數據集選擇：用于量化的數據集要盡可能覆蓋實際應用場景中的數據分布。如果數據集與實際使用的數據差異較大，可能導致量化后的模型精度嚴重下降。比如，在目標檢測任務中，如果數據集中缺少某些類別的樣本，量化后的模型可能對這些類別檢測效果不佳。
量化模式：RKNN 支持多種量化模式，如動態量化和靜態量化等。本文示例中的量化屬于靜態量化，需要提供數據集進行校準。動態量化則不需要額外的校準數據集，但在某些復雜模型上可能精度不如靜態量化。開發者需要根據模型特點和實際需求選擇合適的量化模式。
精度損失評估：量化不可避免會帶來一定的精度損失，在導出模型后，需要通過在測試集上進行推理，對比量化前后模型的準確率、召回率等指標，評估量化效果。如果精度損失過大，可能需要調整量化參數，或者嘗試其他量化策略。

通過以上步驟和解析，我們就完成了從 ONNX 模型到 RKNN 模型的轉換與量化，并在本地驗證了模型推理的效果。下面我給出完整代碼：：

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNNONNX_MODEL = 'yolov5s.onnx'
RKNN_MODEL = 'yolov5s_quant.rknn'
IMG_PATH = './street.jpg'
DATASET = './datasets.txt'QUANTIZE_ON = TrueOBJ_THRESH = 0.25
NMS_THRESH = 0.45
IMG_SIZE = 640CLASSES = ("person", "bicycle", "car", "motorbike ", "aeroplane ", "bus ", "train", "truck ", "boat", "traffic light","fire hydrant", "stop sign ", "parking meter", "bench", "bird", "cat", "dog ", "horse ", "sheep", "cow", "elephant","bear", "zebra ", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite","baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife ","spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza ", "donut", "cake", "chair", "sofa","pottedplant", "bed", "diningtable", "toilet ", "tvmonitor", "laptop	", "mouse	", "remote ", "keyboard ", "cell phone", "microwave ","oven ", "toaster", "sink", "refrigerator ", "book", "clock", "vase", "scissors ", "teddy bear ", "hair drier", "toothbrush ")def sigmoid(x):return 1 / (1 + np.exp(-x))def xywh2xyxy(x):# Convert [x, y, w, h] to [x1, y1, x2, y2]y = np.copy(x)y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left xy[:, 1] = x[:, 1] - x[:, 3] / 2  # top left yy[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right xy[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right yreturn ydef process(input, mask, anchors):anchors = [anchors[i] for i in mask]grid_h, grid_w = map(int, input.shape[0:2])box_confidence = sigmoid(input[..., 4])box_confidence = np.expand_dims(box_confidence, axis=-1)box_class_probs = sigmoid(input[..., 5:])box_xy = sigmoid(input[..., :2])*2 - 0.5col = np.tile(np.arange(0, grid_w), grid_w).reshape(-1, grid_w)row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_h)col = col.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)row = row.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)grid = np.concatenate((col, row), axis=-1)box_xy += gridbox_xy *= int(IMG_SIZE/grid_h)box_wh = pow(sigmoid(input[..., 2:4])*2, 2)box_wh = box_wh * anchorsbox = np.concatenate((box_xy, box_wh), axis=-1)return box, box_confidence, box_class_probsdef filter_boxes(boxes, box_confidences, box_class_probs):"""Filter boxes with box threshold. It's a bit different with origin yolov5 post process!# Argumentsboxes: ndarray, boxes of objects.box_confidences: ndarray, confidences of objects.box_class_probs: ndarray, class_probs of objects.# Returnsboxes: ndarray, filtered boxes.classes: ndarray, classes for boxes.scores: ndarray, scores for boxes."""boxes = boxes.reshape(-1, 4)box_confidences = box_confidences.reshape(-1)box_class_probs = box_class_probs.reshape(-1, box_class_probs.shape[-1])_box_pos = np.where(box_confidences >= OBJ_THRESH)boxes = boxes[_box_pos]box_confidences = box_confidences[_box_pos]box_class_probs = box_class_probs[_box_pos]class_max_score = np.max(box_class_probs, axis=-1)classes = np.argmax(box_class_probs, axis=-1)_class_pos = np.where(class_max_score >= OBJ_THRESH)boxes = boxes[_class_pos]classes = classes[_class_pos]scores = (class_max_score* box_confidences)[_class_pos]return boxes, classes, scoresdef nms_boxes(boxes, scores):"""Suppress non-maximal boxes.# Argumentsboxes: ndarray, boxes of objects.scores: ndarray, scores of objects.# Returnskeep: ndarray, index of effective boxes."""x = boxes[:, 0]y = boxes[:, 1]w = boxes[:, 2] - boxes[:, 0]h = boxes[:, 3] - boxes[:, 1]areas = w * horder = scores.argsort()[::-1]keep = []while order.size > 0:i = order[0]keep.append(i)xx1 = np.maximum(x[i], x[order[1:]])yy1 = np.maximum(y[i], y[order[1:]])xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)inter = w1 * h1ovr = inter / (areas[i] + areas[order[1:]] - inter)inds = np.where(ovr <= NMS_THRESH)[0]order = order[inds + 1]keep = np.array(keep)return keepdef yolov5_post_process(input_data):masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]anchors = [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45],[59, 119], [116, 90], [156, 198], [373, 326]]boxes, classes, scores = [], [], []for input, mask in zip(input_data, masks):b, c, s = process(input, mask, anchors)b, c, s = filter_boxes(b, c, s)boxes.append(b)classes.append(c)scores.append(s)boxes = np.concatenate(boxes)boxes = xywh2xyxy(boxes)classes = np.concatenate(classes)scores = np.concatenate(scores)nboxes, nclasses, nscores = [], [], []for c in set(classes):inds = np.where(classes == c)b = boxes[inds]c = classes[inds]s = scores[inds]keep = nms_boxes(b, s)nboxes.append(b[keep])nclasses.append(c[keep])nscores.append(s[keep])if not nclasses and not nscores:return None, None, Noneboxes = np.concatenate(nboxes)classes = np.concatenate(nclasses)scores = np.concatenate(nscores)return boxes, classes, scoresdef draw(image, boxes, scores, classes):"""Draw the boxes on the image.# Argument:image: original image.boxes: ndarray, boxes of objects.classes: ndarray, classes of objects.scores: ndarray, scores of objects.all_classes: all classes name."""for box, score, cl in zip(boxes, scores, classes):top, left, right, bottom = boxprint('class: {}, score: {}'.format(CLASSES[cl], score))print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom))top = int(top)left = int(left)right = int(right)bottom = int(bottom)cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),(top, left - 6),cv2.FONT_HERSHEY_SIMPLEX,0.6, (0, 0, 255), 2)def letterbox(im, new_shape=(640, 640), color=(0, 0, 0)):# Resize and pad image while meeting stride-multiple constraintsshape = im.shape[:2]  # current shape [height, width]if isinstance(new_shape, int):new_shape = (new_shape, new_shape)# Scale ratio (new / old)r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])# Compute paddingratio = r, r  # width, height ratiosnew_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh paddingdw /= 2  # divide padding into 2 sidesdh /= 2if shape[::-1] != new_unpad:  # resizeim = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))left, right = int(round(dw - 0.1)), int(round(dw + 0.1))im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add borderreturn im, ratio, (dw, dh)if __name__ == '__main__':# Create RKNN objectrknn = RKNN(verbose=True)# pre-process configprint('--> Config model')rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform='rk3588')print('done')# Load ONNX modelprint('--> Loading model')ret = rknn.load_onnx(model=ONNX_MODEL)if ret != 0:print('Load model failed!')exit(ret)print('done')# Build modelprint('--> Building model')ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)if ret != 0:print('Build model failed!')exit(ret)print('done')# Export RKNN modelprint('--> Export rknn model')ret = rknn.export_rknn(RKNN_MODEL)if ret != 0:print('Export rknn model failed!')exit(ret)print('done')# Init runtime environmentprint('--> Init runtime environment')ret = rknn.init_runtime()# ret = rknn.init_runtime('rk3566')if ret != 0:print('Init runtime environment failed!')exit(ret)print('done')# Set inputsimg = cv2.imread(IMG_PATH)# img, ratio, (dw, dh) = letterbox(img, new_shape=(IMG_SIZE, IMG_SIZE))img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))# Inferenceprint('--> Running model')outputs = rknn.inference(inputs=[img])np.save('./onnx_yolov5_0.npy', outputs[0])np.save('./onnx_yolov5_1.npy', outputs[1])np.save('./onnx_yolov5_2.npy', outputs[2])print('done')# post processinput0_data = outputs[0]input1_data = outputs[1]input2_data = outputs[2]input0_data = input0_data.reshape([3, -1]+list(input0_data.shape[-2:]))input1_data = input1_data.reshape([3, -1]+list(input1_data.shape[-2:]))input2_data = input2_data.reshape([3, -1]+list(input2_data.shape[-2:]))input_data = list()input_data.append(np.transpose(input0_data, (2, 3, 0, 1)))input_data.append(np.transpose(input1_data, (2, 3, 0, 1)))input_data.append(np.transpose(input2_data, (2, 3, 0, 1)))boxes, classes, scores = yolov5_post_process(input_data)img_1 = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)if boxes is not None:draw(img_1, boxes, scores, classes)cv2.imwrite('result.png', img_1)rknn.release()