未验证 提交 eb15e9a7 编写于 作者: Z zhupengyang 提交者: GitHub

enhance yolo_box_fuse_pass (#42926)

上级 18323a46
...@@ -199,9 +199,11 @@ void YoloBoxFusePass::ApplyImpl(ir::Graph* graph) const { ...@@ -199,9 +199,11 @@ void YoloBoxFusePass::ApplyImpl(ir::Graph* graph) const {
GET_IR_NODE(nms_out_rois_num); GET_IR_NODE(nms_out_rois_num);
#undef GET_IR_NODE #undef GET_IR_NODE
auto* block = yolo_box0->Op()->Block();
// create yolo_box_head // create yolo_box_head
#define CREATE_YOLO_BOX_HEAD(idx_) \ #define CREATE_YOLO_BOX_HEAD(idx_) \
framework::OpDesc yolo_box_head##idx_##_op_desc; \ framework::OpDesc yolo_box_head##idx_##_op_desc(block); \
yolo_box_head##idx_##_op_desc.SetType("yolo_box_head"); \ yolo_box_head##idx_##_op_desc.SetType("yolo_box_head"); \
yolo_box_head##idx_##_op_desc.SetInput("X", \ yolo_box_head##idx_##_op_desc.SetInput("X", \
{yolo_box##idx_##_in_x->Name()}); \ {yolo_box##idx_##_in_x->Name()}); \
...@@ -222,7 +224,7 @@ void YoloBoxFusePass::ApplyImpl(ir::Graph* graph) const { ...@@ -222,7 +224,7 @@ void YoloBoxFusePass::ApplyImpl(ir::Graph* graph) const {
#undef CREATE_YOLO_BOX_HEAD #undef CREATE_YOLO_BOX_HEAD
// create yolo_box_post // create yolo_box_post
framework::OpDesc yolo_box_post_op_desc; framework::OpDesc yolo_box_post_op_desc(block);
yolo_box_post_op_desc.SetType("yolo_box_post"); yolo_box_post_op_desc.SetType("yolo_box_post");
yolo_box_post_op_desc.SetInput("Boxes0", {yolo_box0_out_boxes->Name()}); yolo_box_post_op_desc.SetInput("Boxes0", {yolo_box0_out_boxes->Name()});
yolo_box_post_op_desc.SetInput("Boxes1", {yolo_box1_out_boxes->Name()}); yolo_box_post_op_desc.SetInput("Boxes1", {yolo_box1_out_boxes->Name()});
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册