[FIX] Fix input/output data size calculation

This patch fixes input/output data size calculation.

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

diff --git a/src/core/ne-model.cc b/src/core/ne-model.cc
index e2bb0e0..0244136 100644 (file)
--- a/src/core/ne-model.cc
+++ b/src/core/ne-model.cc
@@ -241,20 +241,23 @@ Metadata_v3::getInputTensorSize (uint32_t idx, data_layout layout) const
 
   const uint32_t *dims = getInputDims (idx);
   uint32_t tensor_size = getInputElemSize (idx);
+  uint32_t depth = dims[3];
+  bool image;
 
   for (uint32_t rank_idx = 0; rank_idx < MAX_RANK; rank_idx++)
     tensor_size *= dims[rank_idx];
 
+  image = (depth == 1 || depth == 3);
+
   /** special handling for TRIV2 */
-  if (layout == DATA_LAYOUT_TRIV2 && dims[3] != 3 &&
-      dims[3] % DATA_GRANULARITY != 0) {
-    uint32_t depth;
+  if (layout == DATA_LAYOUT_TRIV2 && depth % DATA_GRANULARITY != 0 && !image) {
+    uint32_t new_depth;
 
-    depth = (dims[3] + DATA_GRANULARITY - 1) >> DATA_GRANULARITY_SHIFT;
-    depth = depth * DATA_GRANULARITY;
+    new_depth = (depth + DATA_GRANULARITY - 1) >> DATA_GRANULARITY_SHIFT;
+    new_depth = new_depth * DATA_GRANULARITY;
 
-    tensor_size /= dims[3];
-    tensor_size *= depth;
+    tensor_size /= depth;
+    tensor_size *= new_depth;
   }
 
   return tensor_size;
@@ -271,20 +274,20 @@ Metadata_v3::getOutputTensorSize (uint32_t idx, data_layout layout) const
 
   const uint32_t *dims = getOutputDims (idx);
   uint32_t tensor_size = getOutputElemSize (idx);
+  uint32_t depth = dims[3];
 
   for (uint32_t rank_idx = 0; rank_idx < MAX_RANK; rank_idx++)
     tensor_size *= dims[rank_idx];
 
   /** special handling for TRIV2 */
-  if (layout == DATA_LAYOUT_TRIV2 && dims[3] != 3 &&
-      dims[3] % DATA_GRANULARITY != 0) {
-    uint32_t depth;
+  if (layout == DATA_LAYOUT_TRIV2 && depth % DATA_GRANULARITY != 0) {
+    uint32_t new_depth;
 
-    depth = (dims[3] + DATA_GRANULARITY - 1) >> DATA_GRANULARITY_SHIFT;
-    depth = depth * DATA_GRANULARITY;
+    new_depth = (depth + DATA_GRANULARITY - 1) >> DATA_GRANULARITY_SHIFT;
+    new_depth = new_depth * DATA_GRANULARITY;
 
-    tensor_size /= dims[3];
-    tensor_size *= depth;
+    tensor_size /= depth;
+    tensor_size *= new_depth;
   }
 
   return tensor_size;