关于Detectron2的GenerateRCNN模型转RKNN问题

大佬们好：
我要将一个使用Detectron2框架训练的GenerateRCNN模型转换成rknn模型，遇到困难。大佬们可以帮忙看一下吗？
报错如图：

我的转换代码：

import torch
import os

from rknn.api import RKNN

from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import detection_utils, build_detection_test_loader
from detectron2.export import TracingAdapter
from detectron2.modeling import build_model
from detectron2.utils.file_io import PathManager
import detectron2.data.transforms as ts



def setup_cfg(cfg_path, m_path):
    """初始化配置文件"""
    cfg = get_cfg()
    cfg.DATALOADER.NUM_WORKERS = 0  # 禁用多进程，避免CUDA上下文问题
    cfg.merge_from_file(cfg_path)
    cfg.MODEL.WEIGHTS = os.path.join(m_path, "model_final.pth")
    cfg.MODEL.DEVICE = "cpu"  # 确保在CPU上导出，避免设备不兼容

    cfg.freeze()
    return cfg

def get_sample_inputs(cfg, sample_image_path=None):
    """生成模型导出所需的样本输入"""
    if sample_image_path:
        # 从指定图片生成输入
        original_image = detection_utils.read_image(sample_image_path, format=cfg.INPUT.FORMAT)
        # 应用与推理时相同的预处理
        aug = ts.ResizeShortestEdge(
            [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST],
            cfg.INPUT.MAX_SIZE_TEST
        )
        height, width = original_image.shape[:2]
        image = aug.get_transform(original_image).apply_image(original_image)
        image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))  # 转换为CHW格式
        return [{"image": image, "height": height, "width": width}]
    else:
        # 从测试数据集获取第一个批次作为输入
        data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0])
        return next(iter(data_loader))

def export_generalized_rcnn_to_onnx(cfg, model, sample_inputs, output_path):
    """将GeneralizedRCNN模型导出为ONNX格式"""
    # 提取图像输入（移除其他无用键）
    image = sample_inputs[0]["image"]
    inputs = [{"image": image}]

    # 定义GeneralizedRCNN的推理函数（用于追踪）
    def rcnn_inference(model, inputs):
        # 禁用后处理，保留原始输出用于导出
        instances = model.inference(inputs, do_postprocess=False)[0]
        return [{"instances": instances}]

    # 创建追踪适配器，处理模型输入输出格式
    traceable_model = TracingAdapter(model, inputs, rcnn_inference)

    # 导出为ONNX
    with PathManager.open(output_path, "wb") as f:
        torch.onnx.export(
            traceable_model,
            (image,),  # 输入张量
            f,
            input_names=["input_image"],
            opset_version=11,  # 使用稳定的ONNX算子集版本
            do_constant_folding=True,  # 启用常量折叠优化
            verbose=False  # 禁用详细输出
        )

    print(f"ONNX模型已保存至: {output_path}")
    print(f"输入格式: {traceable_model.inputs_schema}")
    print(f"输出格式: {traceable_model.outputs_schema}")

def onnx2rknn(onnx_path, rknn_path, target_platform='rk3588'):
    # 初始化RKNN对象
    rknn = RKNN(verbose=True)

    # 配置模型输入（根据实际模型的输入尺寸和预处理方式调整）
    # 例如：输入尺寸为(3, 224, 224)，均值和标准差用于归一化
    args = {
        'mean_values': [[123.675, 116.28, 103.53]],  # RGB通道均值（根据模型训练时的预处理设置）
        'std_values': [[58.395, 57.12, 57.375]],     # RGB通道标准差
        'target_platform': target_platform
    }
    rknn.config(
        mean_values=[[103.53, 116.28, 123.675]],
        std_values=[[1.0, 1.0, 1.0]],
        target_platform='rk3588')
    # rknn.config(args)
    input_size_list = [[1, 3, 800, 800]]  # 固定为640x640
    # 加载ONNX模型
    print('--> Loading ONNX model...')
    ret = rknn.load_onnx(model=onnx_path, inputs=['input_image'], input_size_list= input_size_list)
    if ret != 0:
        print('Load ONNX model failed!')
        exit(ret)

    # 构建RKNN模型（会进行算子转换和优化）
    print('--> Building model...')
    ret = rknn.build(do_quantization=False)  # 初始可先不量化，成功后再尝试量化
    if ret != 0:
        print('Build model failed!')
        exit(ret)

    # 导出RKNN模型
    print('--> Exporting RKNN model...')
    ret = rknn.export_rknn(rknn_path)
    if ret != 0:
        print('Export RKNN model failed!')
        exit(ret)

    # 释放资源
    rknn.release()
    print('ONNX to RKNN conversion completed!')


if __name__ == "__main__":
    # 配置路径（请替换为你的实际路径）
    config_path = './detectron2/models/Config/config.yaml'
    model_path = "./detectron2/models"  # 确保该路径下有模型权重文件

    onnx_path = "mask_rcnn.onnx"
    onnx_path_s = "mask_rcnn_simplified.onnx"

    # 加载配置
    # cfg = setup_cfg(config_path, model_path)
    #
    # # 构建并加载模型
    # model = build_model(cfg)
    # DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS)  # 加载权重
    # model.eval()  # 设置为评估模式
    #
    # sample_inputs = get_sample_inputs(cfg,  "./test.png")
    # sample_inputs1 = torch.randn(3, 800, 800)
    # sample_inputs2 = [{"image": sample_inputs1, "height": 16, "width": 32}]
    # # 导出ONNX
    # export_generalized_rcnn_to_onnx(cfg, model, sample_inputs2, onnx_path)
    onnx2rknn(onnx_path_s, "mask_rcnn.rknn", target_platform='rk3588')