[PT FE] Torchvision NMS can accept negative scores (#19826)

This commit is contained in:
Maxim Vafin
2023-09-14 11:07:24 +02:00
committed by GitHub
parent 16adb01810
commit 1a950f9e8d
2 changed files with 6 additions and 3 deletions

View File

@@ -34,8 +34,10 @@ OutputVector translate_nms(const NodeContext& context) {
context.mark_node(v0::Constant::create(element::i32, Shape{1}, {std::numeric_limits<int32_t>::max()}));
auto iou_threshold = context.get_input(2);
auto nms_out =
context.mark_node(std::make_shared<v9::NonMaxSuppression>(boxes, scores, max_output_per_class, iou_threshold));
auto score_threshold =
context.mark_node(v0::Constant::create(element::f32, Shape{}, {std::numeric_limits<float>::lowest()}));
auto nms_out = context.mark_node(
std::make_shared<v9::NonMaxSuppression>(boxes, scores, max_output_per_class, iou_threshold, score_threshold));
auto select = context.mark_node(std::make_shared<v8::Gather>(nms_out, const_2, const_1));
return {context.mark_node(std::make_shared<v0::Squeeze>(select, const_1))};

View File

@@ -15,7 +15,8 @@ class TestNms(PytorchLayerTest):
# PyTorch requires that boxes are in (x1, y1, x2, y2) format, where 0<=x1<x2 and 0<=y1<y2
boxes = np.array([[np.random.uniform(1, 3), np.random.uniform(2, 6),
np.random.uniform(4, 6), np.random.uniform(7, 9)] for _ in range(self.boxes_num)]).astype(np.float32)
scores = np.abs(np.random.randn(self.boxes_num).astype(np.float32))
# scores can be negative
scores = np.random.randn(self.boxes_num).astype(np.float32)
return (boxes, scores)
def create_model(self):