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

enhance yolo_box_fuse_pass (#42926)

上级 18323a46
......@@ -199,9 +199,11 @@ void YoloBoxFusePass::ApplyImpl(ir::Graph* graph) const {
GET_IR_NODE(nms_out_rois_num);
#undef GET_IR_NODE
auto* block = yolo_box0->Op()->Block();
// create yolo_box_head
#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.SetInput("X", \
{yolo_box##idx_##_in_x->Name()}); \
......@@ -222,7 +224,7 @@ void YoloBoxFusePass::ApplyImpl(ir::Graph* graph) const {
#undef CREATE_YOLO_BOX_HEAD
// 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.SetInput("Boxes0", {yolo_box0_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.
先完成此消息的编辑!
想要评论请 注册