海思hi3519dv500,hi3516dv500移植yolov8-模型處理

LittroInno發表於2024-03-05

  上一節yolov8的訓練已經完成了,現在要開始做模型的轉換了,這裡和yolov7方式相似,但是有一些差異,尤其是yolov7的不帶NMS部分的輸出順序和yolov8的輸出順序與格式是有差異的。

   首先還是要自己手動加入rpn_op,這裡包含了filter,sort,nms部分。

我們一個一個看,首先filter.py裡 要先增加個Filter的自定義層。 

op_source = """
        #include <torch/script.h>
        torch::Tensor Filter(torch::Tensor in, int64_t topK, double m_lowScoreThresh)
        {
            return in;
        }
        static auto registry =
            torch::RegisterOperators("custom_ops::Filter", &Filter);
        """

torch.utils.cpp_extension.load_inline(
    name='Filter',
    cpp_sources=op_source,
    is_python_module=False,
    verbose=True,
    )

@parse_args("v", "i", "f")
def symbolic_filte

相關文章