Convert tensor from NHWC (TFLite) to HWCN (#353)

Add a function that creates an IR tensor from a TFLite tensor

Note: this code is subject to change if the way IR stores tensor
data and type changes.

Signed-off-by: Dmitry Mozolev d.mozolev@samsung.com

diff --git a/contrib/nnc/libs/frontend/tflite/include/tflite_ir_visitor.h b/contrib/nnc/libs/frontend/tflite/include/tflite_ir_visitor.h
index 06a48aa..78999c4 100644 (file)
--- a/contrib/nnc/libs/frontend/tflite/include/tflite_ir_visitor.h
+++ b/contrib/nnc/libs/frontend/tflite/include/tflite_ir_visitor.h
@@ -58,6 +58,7 @@ private:
   std::map<int, INode::Ref> opsForTensorsTheyOutput;
 
   std::shared_ptr<IrTensor> createTensor(const Tensor *t, const Buffer *b);
+  std::shared_ptr<IrTensor> convertTensorForConv(std::shared_ptr<IrTensor>);
 };
 
 } // namespace tflite

diff --git a/contrib/nnc/libs/frontend/tflite/src/tflite_ir_visitor.cpp b/contrib/nnc/libs/frontend/tflite/src/tflite_ir_visitor.cpp
index 90907c6..80071ec 100644 (file)
--- a/contrib/nnc/libs/frontend/tflite/src/tflite_ir_visitor.cpp
+++ b/contrib/nnc/libs/frontend/tflite/src/tflite_ir_visitor.cpp
@@ -4,6 +4,8 @@
 #include "schema_v3.h"
 #include "PluginException.h"
 #include "nncc/core/ADT/tensor/Shape.h"
+#include "nncc/core/ADT/tensor/Index.h"
+#include "nncc/core/ADT/tensor/IndexRange.h"
 #include "nnc/core/IR/model/operations/variable_op.h"
 
 #include "shape_helper.h"
@@ -18,6 +20,8 @@ namespace frontend
 namespace tflite
 {
 
+using nncc::core::ADT::tensor::Index;
+using nncc::core::ADT::tensor::IndexRange;
 using VariableOp = nncc::contrib::core::IR::model::ops::VariableOp;
 using nncc::core::ADT::tensor::Shape;
 
@@ -108,6 +112,39 @@ std::shared_ptr<IrTensor> IrVisitor::createTensor(const Tensor *t, const Buffer
   return std::make_shared<IrTensor>(tensorShape, tensorBufferCopy, type, elementSize);
 }
 
+std::shared_ptr<IrTensor> IrVisitor::convertTensorForConv(std::shared_ptr<IrTensor> tensor)
+{
+  Shape inShape = tensor->getShape();
+  assert(inShape.rank() == 4);
+
+  // Changes dimension indices [0, 1, 2, 3] to [1, 2, 3, 0].
+  // This is needed because in TFLite convolution weights are stored as NHWC, and we use HWCN.
+  Shape targetShape{inShape.dim(1), inShape.dim(2), inShape.dim(3), inShape.dim(0)};
+
+  int size = targetShape.dim(0) * targetShape.dim(1) * targetShape.dim(2) * targetShape.dim(3);
+
+  std::shared_ptr<char> convertedTensorData(new char[size * tensor->getElementSize()],
+                                            [](char *d) { delete[] d; });
+
+  auto convertedTensor = std::make_shared<IrTensor>(targetShape, convertedTensorData,
+                                                    tensor->getDataType(),
+                                                    tensor->getElementSize());
+
+  // Swaps two elements in the initial and new tensors
+  // according to the correct tensor dimension order
+  auto swap = [&convertedTensor, &tensor](const Index &idx) {
+    Index targetIndex{idx.at(1), idx.at(2), idx.at(3), idx.at(0)};
+
+    // Copy element_size bytes which constitute one number
+    std::memcpy(convertedTensor->at(targetIndex), tensor->at(idx),
+                convertedTensor->getElementSize());
+  };
+
+  IndexRange(tensor->getShape()).iterate() << swap;
+
+  return convertedTensor;
+}
+
 Graph *IrVisitor::getGraph() { return graph; }
 
 void IrVisitor::setGraphOutputs()