提交 843a95a3 编写于 作者: C Channingss

fix bug of multiclass_nms when attr:keep_top_k==-1

上级 47a244f6
...@@ -72,7 +72,7 @@ def multiclass_nms(op, block): ...@@ -72,7 +72,7 @@ def multiclass_nms(op, block):
dims=(), dims=(),
vals=[float(attrs['nms_threshold'])])) vals=[float(attrs['nms_threshold'])]))
boxes_num = block.var( outputs['Out'][0]).shape[0] boxes_num = block.var(outputs['Out'][0]).shape[0]
top_k_value = np.int64(boxes_num if attrs['keep_top_k'] == -1 else attrs['keep_top_k']) top_k_value = np.int64(boxes_num if attrs['keep_top_k'] == -1 else attrs['keep_top_k'])
node_keep_top_k = onnx.helper.make_node( node_keep_top_k = onnx.helper.make_node(
'Constant', 'Constant',
......
...@@ -72,7 +72,7 @@ def multiclass_nms(op, block): ...@@ -72,7 +72,7 @@ def multiclass_nms(op, block):
dims=(), dims=(),
vals=[float(attrs['nms_threshold'])])) vals=[float(attrs['nms_threshold'])]))
boxes_num = block.var( outputs['Out'][0]).shape[0] boxes_num = block.var(outputs['Out'][0]).shape[0]
top_k_value = np.int64(boxes_num if attrs['keep_top_k'] == -1 else attrs['keep_top_k']) top_k_value = np.int64(boxes_num if attrs['keep_top_k'] == -1 else attrs['keep_top_k'])
node_keep_top_k = onnx.helper.make_node( node_keep_top_k = onnx.helper.make_node(
'Constant', 'Constant',
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册