COMPMID-1097: Port mobilenet to NHWC

Change-Id: I789065bfa0d4ef133388e1904c5caf31e450f80f
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/129495
Tested-by: Jenkins <bsgcomp@arm.com>
Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>

3 files changed, 104 insertions, 23 deletions

diff --git a/utils/GraphUtils.cpp b/utils/GraphUtils.cpp
index 7912fd6b7d..145e44950b 100644
--- a/utils/GraphUtils.cpp
+++ b/utils/GraphUtils.cpp
@@ -24,6 +24,8 @@
 
 #include "utils/GraphUtils.h"
 
+#include "arm_compute/core/Helpers.h"
+#include "arm_compute/core/Types.h"
 #include "arm_compute/runtime/SubTensor.h"
 #include "utils/Utils.h"
 
@@ -31,6 +33,27 @@
 
 using namespace arm_compute::graph_utils;
 
+namespace
+{
+std::pair<arm_compute::TensorShape, arm_compute::PermutationVector> compute_permutation_paramaters(const arm_compute::TensorShape &shape,
+                                                                                                   arm_compute::DataLayout data_layout)
+{
+    // Set permutation parameters if needed
+    arm_compute::TensorShape       permuted_shape = shape;
+    arm_compute::PermutationVector perm;
+    // Permute only if num_dimensions greater than 2
+    if(shape.num_dimensions() > 2)
+    {
+        perm = (data_layout == arm_compute::DataLayout::NHWC) ? arm_compute::PermutationVector(2U, 0U, 1U) : arm_compute::PermutationVector(1U, 2U, 0U);
+
+        arm_compute::PermutationVector perm_shape = (data_layout == arm_compute::DataLayout::NCHW) ? arm_compute::PermutationVector(2U, 0U, 1U) : arm_compute::PermutationVector(1U, 2U, 0U);
+        arm_compute::permute(permuted_shape, perm_shape);
+    }
+
+    return std::make_pair(permuted_shape, perm);
+}
+} // namespace
+
 void TFPreproccessor::preprocess(ITensor &tensor)
 {
     Window window;
@@ -118,8 +141,15 @@ bool PPMAccessor::access_tensor(ITensor &tensor)
     // Open PPM file
     ppm.open(_ppm_path);
 
-    ARM_COMPUTE_ERROR_ON_MSG(ppm.width() != tensor.info()->dimension(0) || ppm.height() != tensor.info()->dimension(1),
-                             "Failed to load image file: dimensions [%d,%d] not correct, expected [%d,%d].", ppm.width(), ppm.height(), tensor.info()->dimension(0), tensor.info()->dimension(1));
+    // Get permutated shape and permutation parameters
+    TensorShape                    permuted_shape = tensor.info()->tensor_shape();
+    arm_compute::PermutationVector perm;
+    if(tensor.info()->data_layout() != DataLayout::NCHW)
+    {
+        std::tie(permuted_shape, perm) = compute_permutation_paramaters(tensor.info()->tensor_shape(), tensor.info()->data_layout());
+    }
+    ARM_COMPUTE_ERROR_ON_MSG(ppm.width() != permuted_shape.x() || ppm.height() != permuted_shape.y(),
+                             "Failed to load image file: dimensions [%d,%d] not correct, expected [%d,%d].", ppm.width(), ppm.height(), permuted_shape.x(), permuted_shape.y());
 
     // Fill the tensor with the PPM content (BGR)
     ppm.fill_planar_tensor(tensor, _bgr);
@@ -320,8 +350,8 @@ bool RandomAccessor::access_tensor(ITensor &tensor)
     return true;
 }
 
-NumPyBinLoader::NumPyBinLoader(std::string filename)
-    : _filename(std::move(filename))
+NumPyBinLoader::NumPyBinLoader(std::string filename, DataLayout file_layout)
+    : _filename(std::move(filename)), _file_layout(file_layout)
 {
 }
 
@@ -366,30 +396,57 @@ bool NumPyBinLoader::access_tensor(ITensor &tensor)
         }
     }
 
+    bool are_layouts_different = (_file_layout != tensor.info()->data_layout());
+
     // Validate tensor ranks
     ARM_COMPUTE_ERROR_ON_MSG(shape.size() != tensor_shape.num_dimensions(), "Tensor ranks mismatch");
 
+    // Set permutation parameters if needed
+    TensorShape                    permuted_shape = tensor_shape;
+    arm_compute::PermutationVector perm;
+    if(are_layouts_different)
+    {
+        std::tie(permuted_shape, perm) = compute_permutation_paramaters(tensor_shape, tensor.info()->data_layout());
+    }
+
     // Validate shapes
     for(size_t i = 0; i < shape.size(); ++i)
     {
-        ARM_COMPUTE_ERROR_ON_MSG(tensor_shape[i] != shape[i], "Tensor dimensions mismatch");
+        ARM_COMPUTE_ERROR_ON_MSG(permuted_shape[i] != shape[i], "Tensor dimensions mismatch");
     }
 
-    // Read data
-    if(tensor.info()->padding().empty() && (dynamic_cast<SubTensor *>(&tensor) == nullptr))
+    // Validate shapes and copy tensor
+    if(!are_layouts_different || perm.num_dimensions() <= 2)
     {
-        // If tensor has no padding read directly from stream.
-        stream.read(reinterpret_cast<char *>(tensor.buffer()), tensor.info()->total_size());
+        // Read data
+        if(tensor.info()->padding().empty() && (dynamic_cast<SubTensor *>(&tensor) == nullptr))
+        {
+            // If tensor has no padding read directly from stream.
+            stream.read(reinterpret_cast<char *>(tensor.buffer()), tensor.info()->total_size());
+        }
+        else
+        {
+            // If tensor has padding accessing tensor elements through execution window.
+            Window window;
+            window.use_tensor_dimensions(tensor_shape);
+
+            execute_window_loop(window, [&](const Coordinates & id)
+            {
+                stream.read(reinterpret_cast<char *>(tensor.ptr_to_element(id)), tensor.info()->element_size());
+            });
+        }
     }
     else
     {
         // If tensor has padding accessing tensor elements through execution window.
         Window window;
-        window.use_tensor_dimensions(tensor_shape);
+        window.use_tensor_dimensions(permuted_shape);
 
         execute_window_loop(window, [&](const Coordinates & id)
         {
-            stream.read(reinterpret_cast<char *>(tensor.ptr_to_element(id)), tensor.info()->element_size());
+            Coordinates coords(id);
+            arm_compute::permute(coords, perm);
+            stream.read(reinterpret_cast<char *>(tensor.ptr_to_element(coords)), tensor.info()->element_size());
         });
     }
     return true;
diff --git a/utils/GraphUtils.h b/utils/GraphUtils.h
index 11f1e0590a..a8507b1ac7 100644
--- a/utils/GraphUtils.h
+++ b/utils/GraphUtils.h
@@ -201,9 +201,10 @@ class NumPyBinLoader final : public graph::ITensorAccessor
 public:
     /** Default Constructor
      *
-     * @param filename Binary file name
+     * @param[in] filename    Binary file name
+     * @param[in] file_layout (Optional) Layout of the numpy tensor data. Defaults to NCHW
      */
-    NumPyBinLoader(std::string filename);
+    NumPyBinLoader(std::string filename, DataLayout file_layout = DataLayout::NCHW);
     /** Allows instances to move constructed */
     NumPyBinLoader(NumPyBinLoader &&) = default;
 
@@ -212,6 +213,7 @@ public:
 
 private:
     const std::string _filename;
+    const DataLayout  _file_layout;
 };
 
 /** Generates appropriate random accessor
@@ -231,12 +233,15 @@ inline std::unique_ptr<graph::ITensorAccessor> get_random_accessor(PixelValue lo
  *
  * @note If path is empty will generate a DummyAccessor else will generate a NumPyBinLoader
  *
- * @param[in] path      Path to the data files
- * @param[in] data_file Relative path to the data files from path
+ * @param[in] path        Path to the data files
+ * @param[in] data_file   Relative path to the data files from path
+ * @param[in] file_layout (Optional) Layout of file. Defaults to NCHW
  *
  * @return An appropriate tensor accessor
  */
-inline std::unique_ptr<graph::ITensorAccessor> get_weights_accessor(const std::string &path, const std::string &data_file)
+inline std::unique_ptr<graph::ITensorAccessor> get_weights_accessor(const std::string &path,
+                                                                    const std::string &data_file,
+                                                                    DataLayout         file_layout = DataLayout::NCHW)
 {
     if(path.empty())
     {
@@ -244,7 +249,7 @@ inline std::unique_ptr<graph::ITensorAccessor> get_weights_accessor(const std::s
     }
     else
     {
-        return arm_compute::support::cpp14::make_unique<NumPyBinLoader>(path + data_file);
+        return arm_compute::support::cpp14::make_unique<NumPyBinLoader>(path + data_file, file_layout);
     }
 }
 
diff --git a/utils/Utils.h b/utils/Utils.h
index 6241562a28..cadba3a088 100644
--- a/utils/Utils.h
+++ b/utils/Utils.h
@@ -406,7 +406,14 @@ public:
     {
         ARM_COMPUTE_ERROR_ON(!is_open());
         ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(&tensor, 1, DataType::U8, DataType::F32);
-        ARM_COMPUTE_ERROR_ON(tensor.info()->dimension(0) != _width || tensor.info()->dimension(1) != _height || tensor.info()->dimension(2) != 3);
+
+        const DataLayout  data_layout  = tensor.info()->data_layout();
+        const TensorShape tensor_shape = tensor.info()->tensor_shape();
+
+        ARM_COMPUTE_UNUSED(tensor_shape);
+        ARM_COMPUTE_ERROR_ON(tensor_shape[get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH)] != _width);
+        ARM_COMPUTE_ERROR_ON(tensor_shape[get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT)] != _height);
+        ARM_COMPUTE_ERROR_ON(tensor_shape[get_data_layout_dimension_index(data_layout, DataLayoutDimension::CHANNEL)] != 3);
 
         try
         {
@@ -423,11 +430,25 @@ public:
                                      "Not enough data in file");
             ARM_COMPUTE_UNUSED(end_position);
 
+            // Stride across channels
+            size_t stride_z = 0;
+
             // Iterate through every pixel of the image
             arm_compute::Window window;
-            window.set(arm_compute::Window::DimX, arm_compute::Window::Dimension(0, _width, 1));
-            window.set(arm_compute::Window::DimY, arm_compute::Window::Dimension(0, _height, 1));
-            window.set(arm_compute::Window::DimZ, arm_compute::Window::Dimension(0, 1, 1));
+            if(data_layout == DataLayout::NCHW)
+            {
+                window.set(arm_compute::Window::DimX, arm_compute::Window::Dimension(0, _width, 1));
+                window.set(arm_compute::Window::DimY, arm_compute::Window::Dimension(0, _height, 1));
+                window.set(arm_compute::Window::DimZ, arm_compute::Window::Dimension(0, 1, 1));
+                stride_z = tensor.info()->strides_in_bytes()[2];
+            }
+            else
+            {
+                window.set(arm_compute::Window::DimX, arm_compute::Window::Dimension(0, 1, 1));
+                window.set(arm_compute::Window::DimY, arm_compute::Window::Dimension(0, _width, 1));
+                window.set(arm_compute::Window::DimZ, arm_compute::Window::Dimension(0, _height, 1));
+                stride_z = tensor.info()->strides_in_bytes()[0];
+            }
 
             arm_compute::Iterator out(&tensor, window);
 
@@ -435,8 +456,6 @@ public:
             unsigned char green = 0;
             unsigned char blue  = 0;
 
-            size_t stride_z = tensor.info()->strides_in_bytes()[2];
-
             arm_compute::execute_window_loop(window, [&](const arm_compute::Coordinates & id)
             {
                 red   = _fs.get();