[Model] Fix tensor size calculation

This patch fixes tensor size calculation.
The special handling is applied NONE/TRIV2 layout only.

Signed-off-by: Dongju Chae <dongju.chae@samsung.com>

diff --git a/src/core/ne-handler.cc b/src/core/ne-handler.cc
index 2278c20..ca035cc 100644 (file)
--- a/src/core/ne-handler.cc
+++ b/src/core/ne-handler.cc
@@ -1079,9 +1079,11 @@ TrinityVision2::getTensorSize (const Model *model, bool input, uint32_t index,
     return -EINVAL;
 
   if (input)
-    *size = meta->getInputTensorSize (index, DATA_LAYOUT_TRIV2);
+    *size =
+        meta->getInputTensorSize (index, meta->getInputSegmentLayout (index));
   else
-    *size = meta->getOutputTensorSize (index, DATA_LAYOUT_TRIV2);
+    *size =
+        meta->getOutputTensorSize (index, meta->getOutputSegmentLayout (index));
 
   return 0;
 }

diff --git a/src/core/ne-model.cc b/src/core/ne-model.cc
index 92328c6..287edd7 100644 (file)
--- a/src/core/ne-model.cc
+++ b/src/core/ne-model.cc
@@ -248,36 +248,38 @@ Metadata_v3::getInputTensorSize (uint32_t idx, data_layout layout) const {
   const uint32_t *dims = getInputDims (idx);
   uint32_t elem_size = getInputElemSize (idx);
   uint32_t tensor_size = elem_size;
-  uint32_t depth = dims[3];
-  bool image;
 
   for (uint32_t rank_idx = 0; rank_idx < MAX_RANK; rank_idx++)
     tensor_size *= dims[rank_idx];
 
-  uint32_t granularity = DATA_GRANULARITY;
-  uint32_t shift = DATA_GRANULARITY_SHIFT;
-  uint32_t div = 1;
+  if (layout == DATA_LAYOUT_NONE || layout == DATA_LAYOUT_TRIV2) {
+    uint32_t granularity = DATA_GRANULARITY;
+    uint32_t shift = DATA_GRANULARITY_SHIFT;
+    uint32_t depth = dims[3];
+    uint32_t div = 1;
+    bool image;
 
-  if (getTops () == 2) {
-    div *= 2;
-    shift--;
-  }
-  if (elem_size == 2) {
-    div *= 2;
-    shift--;
-  }
-  granularity /= div;
-  image = (depth == 1 || depth == 3);
+    if (getTops () == 2) {
+      div *= 2;
+      shift--;
+    }
+    if (elem_size == 2) {
+      div *= 2;
+      shift--;
+    }
+    granularity /= div;
+    image = (depth == 1 || depth == 3);
 
-  /** special handling for TRIV2 */
-  if (layout == DATA_LAYOUT_TRIV2 && depth % granularity != 0 && !image) {
-    uint32_t new_depth;
+    /** special handling for TRIV2 */
+    if (layout == DATA_LAYOUT_TRIV2 && depth % granularity != 0 && !image) {
+      uint32_t new_depth;
 
-    new_depth = (depth + granularity - 1) >> shift;
-    new_depth = new_depth * granularity;
+      new_depth = (depth + granularity - 1) >> shift;
+      new_depth = new_depth * granularity;
 
-    tensor_size /= depth;
-    tensor_size *= new_depth;
+      tensor_size /= depth;
+      tensor_size *= new_depth;
+    }
   }
 
   return tensor_size;
@@ -294,34 +296,36 @@ Metadata_v3::getOutputTensorSize (uint32_t idx, data_layout layout) const {
   const uint32_t *dims = getOutputDims (idx);
   uint32_t elem_size = getOutputElemSize (idx);
   uint32_t tensor_size = elem_size;
-  uint32_t depth = dims[3];
 
   for (uint32_t rank_idx = 0; rank_idx < MAX_RANK; rank_idx++)
     tensor_size *= dims[rank_idx];
 
-  uint32_t granularity = DATA_GRANULARITY;
-  uint32_t shift = DATA_GRANULARITY_SHIFT;
-  uint32_t div = 1;
+  if (layout == DATA_LAYOUT_NONE || layout == DATA_LAYOUT_TRIV2) {
+    uint32_t granularity = DATA_GRANULARITY;
+    uint32_t shift = DATA_GRANULARITY_SHIFT;
+    uint32_t depth = dims[3];
+    uint32_t div = 1;
 
-  if (getTops () == 2) {
-    div *= 2;
-    shift--;
-  }
-  if (elem_size == 2) {
-    div *= 2;
-    shift--;
-  }
-  granularity /= div;
+    if (getTops () == 2) {
+      div *= 2;
+      shift--;
+    }
+    if (elem_size == 2) {
+      div *= 2;
+      shift--;
+    }
+    granularity /= div;
 
-  /** special handling for TRIV2 */
-  if (layout == DATA_LAYOUT_TRIV2 && depth % granularity != 0) {
-    uint32_t new_depth;
+    /** special handling for TRIV2 */
+    if (layout == DATA_LAYOUT_TRIV2 && depth % granularity != 0) {
+      uint32_t new_depth;
 
-    new_depth = (depth + granularity - 1) >> shift;
-    new_depth = new_depth * granularity;
+      new_depth = (depth + granularity - 1) >> shift;
+      new_depth = new_depth * granularity;
 
-    tensor_size /= depth;
-    tensor_size *= new_depth;
+      tensor_size /= depth;
+      tensor_size *= new_depth;
+    }
   }
 
   return tensor_size;