深入探索 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 path
        imgsz=(640, 640),  # image (height, width)
        batch_size=1,  # batch size
        device='cpu',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        include=('torchscript', 'onnx', 'coreml'),  # include formats
        half=False,  # FP16 half-precision export
        inplace=False,  # set YOLOv5 Detect() inplace=True
        train=False,  # model.train() mode
        optimize=False,  # TorchScript: optimize for mobile
        int8=False,  # CoreML/TF INT8 quantization
        dynamic=False,  # ONNX/TF: dynamic axes
        simplify=False,  # ONNX: simplify model
        opset=12,  # ONNX: opset version
        verbose=False,  # TensorRT: verbose log
        workspace=4,  # TensorRT: workspace size (GB)
        topk_per_class=100,  # TF.js NMS: topk per class to keep
        topk_all=100,  # TF.js NMS: topk for all classes to keep
        iou_thres=0.45,  # TF.js NMS: IoU threshold
        conf_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 exports
    imgsz *= 2 if len(imgsz) == 1 else 1  # expand
    file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)

    # Load PyTorch model
    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
    nc, names = model.nc, model.names  # number of classes, class names

    # Input
    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

    # Update model
    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 activations
            if isinstance(m.act, nn.SiLU):
                m.act = SiLU()
        elif isinstance(m, Detect):
            m.inplace = inplace
            m.onnx_dynamic = dynamic
            # m.forward = m.forward_export  # assign forward (optional)

    for _ in range(2):
        y = model(im)  # dry runs
    LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")

    # Exports
    if '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 Exports
    if 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 model
        if pb or tfjs:  # pb prerequisite to tfjs
            export_pb(model, im, file)
        if tflite:
            export_tflite(model, im, file, int8=int8, data=data, ncalib=100)
        if tfjs:
            export_tfjs(model, im, file)

    # Finish
    LOGGER.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 activations
        if isinstance(m.act, nn.SiLU):
            m.act = SiLU()
    elif isinstance(m, Detect):
        m.inplace = inplace
        m.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 export
    try:
        check_requirements(('onnx',))
        import onnx

        LOGGER.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)

        # Checks
        model_onnx = onnx.load(f)  # load onnx model
        onnx.checker.check_model(model_onnx)  # check onnx model
        # LOGGER.info(onnx.helper.printable_graph(model_onnx.graph))  # print

        # Simplify
        if simplify:
            try:
                check_requirements(('onnx-simplifier',))
                import onnxsim

                LOGGER.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 onnxsim

        LOGGER.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 之间的转换本质是计算图的映射过程：

1. 计算图提取：从 TensorFlow 的会话 (Session) 中提取计算图结构和权重参数
2. 算子映射：将 TensorFlow 特定算子转换为 ONNX 标准算子
3. 元数据转换：保留输入输出名称、数据类型等元信息
4. 模型序列化：生成符合 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, QuantType

quantized_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 模型，主要分为以下几个步骤：

1. 环境准备：创建 RKNN 对象，配置预处理参数。
2. 模型加载：将 ONNX 格式的模型文件加载到 RKNN 环境中。
3. 模型构建：这一步包含了关键的模型量化操作，将模型转换为适合 RKNN 推理框架运行的形式。
4. 模型导出：生成最终可在瑞芯微设备上运行的 RKNN 模型文件。
5. 推理验证：在本地初始化运行环境，输入测试图像进行推理，并对结果进行后处理展示。

接下来，我们结合代码逐步深入解析每个步骤。

（二）代码实现详解

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.量化过程中的关键要点与注意事项

1. 数据集选择：用于量化的数据集要尽可能覆盖实际应用场景中的数据分布。如果数据集与实际使用的数据差异较大，可能导致量化后的模型精度严重下降。比如，在目标检测任务中，如果数据集中缺少某些类别的样本，量化后的模型可能对这些类别检测效果不佳。
2. 量化模式：RKNN 支持多种量化模式，如动态量化和静态量化等。本文示例中的量化属于静态量化，需要提供数据集进行校准。动态量化则不需要额外的校准数据集，但在某些复杂模型上可能精度不如静态量化。开发者需要根据模型特点和实际需求选择合适的量化模式。
3. 精度损失评估：量化不可避免会带来一定的精度损失，在导出模型后，需要通过在测试集上进行推理，对比量化前后模型的准确率、召回率等指标，评估量化效果。如果精度损失过大，可能需要调整量化参数，或者尝试其他量化策略。

通过以上步骤和解析，我们就完成了从 ONNX 模型到 RKNN 模型的转换与量化，并在本地验证了模型推理的效果。下面我给出完整代码：：

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN

ONNX_MODEL = 'yolov5s.onnx'
RKNN_MODEL = 'yolov5s_quant.rknn'
IMG_PATH = './street.jpg'
DATASET = './datasets.txt'

QUANTIZE_ON = True

OBJ_THRESH = 0.25
NMS_THRESH = 0.45
IMG_SIZE = 640

CLASSES = ("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 x
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
    return y


def 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.5

    col = 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 += grid
    box_xy *= int(IMG_SIZE/grid_h)

    box_wh = pow(sigmoid(input[..., 2:4])*2, 2)
    box_wh = box_wh * anchors

    box = np.concatenate((box_xy, box_wh), axis=-1)

    return box, box_confidence, box_class_probs


def filter_boxes(boxes, box_confidences, box_class_probs):
    """Filter boxes with box threshold. It's a bit different with origin yolov5 post process!

    # Arguments
        boxes: ndarray, boxes of objects.
        box_confidences: ndarray, confidences of objects.
        box_class_probs: ndarray, class_probs of objects.

    # Returns
        boxes: 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, scores


def nms_boxes(boxes, scores):
    """Suppress non-maximal boxes.

    # Arguments
        boxes: ndarray, boxes of objects.
        scores: ndarray, scores of objects.

    # Returns
        keep: ndarray, index of effective boxes.
    """
    x = boxes[:, 0]
    y = boxes[:, 1]
    w = boxes[:, 2] - boxes[:, 0]
    h = boxes[:, 3] - boxes[:, 1]

    areas = w * h
    order = 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 * h1

        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr <= NMS_THRESH)[0]
        order = order[inds + 1]
    keep = np.array(keep)
    return keep


def 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, None

    boxes = np.concatenate(nboxes)
    classes = np.concatenate(nclasses)
    scores = np.concatenate(nscores)

    return boxes, classes, scores


def 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 = box
        print('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 constraints
    shape = 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 padding
    ratio = r, r  # width, height ratios
    new_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 padding

    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # resize
        im = 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 border
    return im, ratio, (dw, dh)


if __name__ == '__main__':

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

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

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

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

    # 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 - toolkit2 中可能没有直接的支持。例如，模型中使用了一种新的卷积变体算子。这时，我们需要查看 rknn - toolkit2 的文档，看是否有替代的实现方式或者是否可以通过组合现有的算子来实现相同的功能。​

如果实在无法通过现有方式实现，可能需要与瑞芯微的技术支持团队沟通，看是否有进一步的解决方案，比如是否可以通过扩展算子库等方式来支持自定义算子。​

五、模型结构适配​

自定义模型的结构可能与标准模型差异较大，例如可能存在一些复杂的分支结构或者嵌套结构。在转换时，需要仔细检查模型结构，确保其符合 rknn - toolkit2 的转换要求。​

有时候可能需要对模型结构进行一些微调，比如将一些过于复杂的分支结构进行简化或者合并。但在进行结构调整时，一定要确保不会影响模型的性能和功能。这就需要我们对模型的原理和功能有深入的理解，在调整前后对模型进行充分的测试和验证。​

通过以上一系列步骤，我们就完成了从常见模型到 ONNX 格式，再借助 Docker 和 rknn - toolkit2 转换为 RKNN 模型的全过程，包括了应对自定义模型转换时的各种挑战。现在，我们手中的 RKNN 模型就可以在瑞芯微强大的硬件平台上高效运行，为各种边缘计算应用提供有力的支持啦！