onnx done

fcos_core/modeling/backbone/fpn.py

@@ -56,11 +56,12 @@ class FPN(nn.Module):
         ):
             if not inner_block:
                 continue
-            inner_top_down = F.interpolate(last_inner, scale_factor=2, mode="nearest")
+            # inner_top_down = F.interpolate(last_inner, scale_factor=2, mode="nearest")
             inner_lateral = getattr(self, inner_block)(feature)
             # TODO use size instead of scale to make it robust to different sizes
-            # inner_top_down = F.upsample(last_inner, size=inner_lateral.shape[-2:],
-            # mode='bilinear', align_corners=False)
+            inner_top_down = F.upsample(
+                last_inner, size=(int(inner_lateral.shape[-2]) + 1 - 1, int(inner_lateral.shape[-1])), mode='nearest'
+            )
             last_inner = inner_lateral + inner_top_down
             results.insert(0, getattr(self, layer_block)(last_inner))
 

fcos_core/modeling/rpn/fcos/inference.py

@@ -9,6 +9,7 @@ from fcos_core.structures.bounding_box import BoxList
 from fcos_core.structures.boxlist_ops import cat_boxlist
 from fcos_core.structures.boxlist_ops import boxlist_ml_nms
 from fcos_core.structures.boxlist_ops import remove_small_boxes
+from torchvision.ops.boxes import batched_nms
 
 
 class FCOSPostProcessor(torch.nn.Module):
@@ -147,7 +148,13 @@ class FCOSPostProcessor(torch.nn.Module):
         results = []
         for i in range(num_images):
             # multiclass nms
-            result = boxlist_ml_nms(boxlists[i], self.nms_thresh)
+            keep = batched_nms(
+                boxlists[i].bbox,
+                boxlists[i].get_field("scores"),
+                boxlists[i].get_field("labels"),
+                self.nms_thresh
+            )
+            result = boxlists[i][keep]
             number_of_detections = len(result)
 
             # Limit to max_per_image detections **over all classes**

tools/test_fcos_onnx_model.py

@@ -21,6 +21,7 @@ from fcos_core.utils.miscellaneous import mkdir
 from fcos_core.modeling.rpn.fcos.inference import make_fcos_postprocessor
 import caffe2.python.onnx.backend as backend
 import numpy as np
+from fcos_core.structures.image_list import to_image_list
 
 
 class ONNX_FCOS(nn.Module):
@@ -37,9 +38,10 @@ class ONNX_FCOS(nn.Module):
     def forward(self, images):
         outputs = self.onnx_model.run(images.tensors.cpu().numpy())
         outputs = [torch.from_numpy(o).to(self.cfg.MODEL.DEVICE) for o in outputs]
-        logits = outputs[::3]
-        bbox_reg = outputs[1::3]
-        centerness = outputs[2::3]
+        num_outputs = len(outputs) // 3
+        logits = outputs[:num_outputs]
+        bbox_reg = outputs[num_outputs:2 * num_outputs]
+        centerness = outputs[2 * num_outputs:]
 
         locations = self.compute_locations(logits)
         boxes = self.postprocessing(locations, logits, bbox_reg, centerness, images.image_sizes)