赞
踩
上一节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_filter
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。