add maskrcnn inference example (intel-analytics#2944)

* add maskrcnn inference example * meet pr comments * add model download url
dding3 · Oct 28, 2019 · 6ffbc0e · 6ffbc0e
1 parent 448d882
commit 6ffbc0e
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 10 deletions.
diff --git a/dl/src/main/scala/com/intel/analytics/bigdl/dllib/nn/layers/BoxHead.scala b/dl/src/main/scala/com/intel/analytics/bigdl/dllib/nn/layers/BoxHead.scala
@@ -342,6 +342,11 @@ private[nn] class BoxPostProcessor(
       totalDetections += maxDetection
       outBBoxs[Tensor[Float]](i + 1).resize(maxDetection, 4)
       totalROILables(i + 1) = roilabels
+      boxesInImage(i) = maxDetection
+    }
+    // clear others tensors in output
+    for (i <- (boxesInImage.length + 1) to outBBoxs.length()) {
+      outBBoxs.remove[Tensor[Float]](i)
     }
 
     // resize labels and scores
@@ -353,13 +358,15 @@ private[nn] class BoxPostProcessor(
     var labelsOffset = outLabels.storageOffset() - 1
     var scoresOffset = outScores.storageOffset() - 1
     for (i <- 0 to boxesInImage.length - 1) {
-      val roilabels = totalROILables[Array[RoiLabel]](i + 1)
-      val bbox = outBBoxs[Tensor[Float]](i + 1).storage().array()
-      val bboxOffset = outBBoxs[Tensor[Float]](i + 1).storageOffset() - 1
-
-      resultToTensor(roilabels, labels, labelsOffset, bbox, bboxOffset, scores, scoresOffset)
-      labelsOffset += outBBoxs[Tensor[Float]](i + 1).size(1)
-      scoresOffset += outBBoxs[Tensor[Float]](i + 1).size(1)
+      if (boxesInImage(i) > 0) {
+        val roilabels = totalROILables[Array[RoiLabel]](i + 1)
+        val bbox = outBBoxs[Tensor[Float]](i + 1).storage().array()
+        val bboxOffset = outBBoxs[Tensor[Float]](i + 1).storageOffset() - 1
+
+        resultToTensor(roilabels, labels, labelsOffset, bbox, bboxOffset, scores, scoresOffset)
+        labelsOffset += outBBoxs[Tensor[Float]](i + 1).size(1)
+        scoresOffset += outBBoxs[Tensor[Float]](i + 1).size(1)
+      }
     }
 
     output

diff --git a/dl/src/main/scala/com/intel/analytics/bigdl/dllib/nn/layers/Pooler.scala b/dl/src/main/scala/com/intel/analytics/bigdl/dllib/nn/layers/Pooler.scala
@@ -109,10 +109,12 @@ class Pooler[T: ClassTag] (
     } else { // for batch support
       input[Table](2)
     }
+
+    val batchSize = featureMaps.get[Tensor[Float]](1).get.size(1)
     var totalNum = 0
     val num_channels = featureMaps.get[Tensor[T]](1).get.size(2)
     val out = T()
-    for (i <- 0 to roiBatch.length() - 1) {
+    for (i <- 0 to batchSize - 1) {
       val rois = roiBatch[Tensor[T]](i + 1)
 
       val roi_levels = levelMapping(lvl_min, lvl_max, rois)
@@ -147,10 +149,10 @@ class Pooler[T: ClassTag] (
     // merge to one tensor
     output.resize(totalNum, num_channels, resolution, resolution)
     var start = 1
-    for (i <- 0 to roiBatch.length() - 1) {
+    for (i <- 0 to batchSize - 1) {
       val tmp = out[Tensor[T]](i + 1)
       val length = tmp.size(1)
-      output.narrow(1, start, length).copy(tmp)
+      if (length > 0) output.narrow(1, start, length).copy(tmp)
       start += length
     }
 

diff --git a/dl/src/main/scala/com/intel/analytics/bigdl/dllib/nn/layers/RegionProposal.scala b/dl/src/main/scala/com/intel/analytics/bigdl/dllib/nn/layers/RegionProposal.scala
@@ -151,6 +151,10 @@ class RegionProposal(
       // sort
       selectOverAllLevels(selectorRes, postNmsTopN, bboxNumber, output[Tensor[Float]](b))
     }
+    // clear others tensors in output
+    for (i <- (batchSize + 1) to output.length()) {
+     output.remove[Tensor[Float]](i)
+    }
     output
   }