fixed instance norm bug & add its layoutinfer (#593)

Type: Bug fix

Signed-off-by: Chen <jack.chen@verisilicon.com>
Co-authored-by: Chen <jack.chen@verisilicon.com>

include/tim/vx/ops/instancenormalization.h

@@ -30,7 +30,7 @@ namespace vx {
 namespace ops {
 class InstanceNormalization : public BuiltinOp {
  public:
-  InstanceNormalization(Graph* graph, float eps = 1e-5f);
+  InstanceNormalization(Graph* graph, float eps = 1e-5f, DataLayout input_layout = DataLayout::WHCN);
 
   std::shared_ptr<Operation> Clone(std::shared_ptr<Graph>& graph) const override;
 

src/tim/transform/layout_inference.cc

@@ -50,6 +50,7 @@
 #include "ops/stridedslice_layout_inference.h"
 #include "ops/lrn_layout_inference.h"
 #include "ops/l2normalization_layout_inference.h"
+#include "ops/instance_norm_layout_inference.h"
 #include "ops/addn_layout_inference.h"
 #include "ops/gather_layout_inference.h"
 #include "ops/gather_nd_layout_inference.h"
@@ -258,6 +259,7 @@ std::vector<std::shared_ptr<vx::Tensor>> HandleLayoutInfer(
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_STRIDED_SLICE, StridedSlice);
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_LRN2, LRN);
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_L2_NORMALIZE, L2Normalization);
+    REGIST_LAYOUT_INFERENCE(VSI_NN_OP_INSTANCE_NORM, InstanceNorm);
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_ADDN, AddN);
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_PRELU, PRelu);
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_GATHER, Gather);
@@ -275,7 +277,7 @@ std::vector<std::shared_ptr<vx::Tensor>> HandleLayoutInfer(
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_EXPAND_BROADCAST, Broadcast);
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_UNIDIRECTIONAL_SEQUENCE_RNN, UnidirectionalRnn);
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_BIDIRECTIONAL_SEQUENCE_RNN, BidirectionalRnn);
-#ifdef ENABLE_TENSOR_CACHE
+#ifdef VSI_FEAT_OP_CUSTOM_TINY_YOLOV4_POSTPROCESS
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_CUSTOM_TINY_YOLOV4_POSTPROCESS, Yolov4);
 #endif
     REGIST_LOGICAL_LAYOUT_INFERENCE(VSI_NN_OP_LOGICAL_OPS);

src/tim/transform/layout_inference_test.cc

@@ -2,6 +2,7 @@
 #include "tim/vx/graph.h"
 #include "tim/vx/ops.h"
 #include "tim/transform/layout_inference.h"
+#include "test_utils.h"
 
 #include "gtest/gtest.h"
 
@@ -220,4 +221,54 @@ TEST(FC, share_const_tensor) {
   std::vector<float> output(golden.size());
   EXPECT_TRUE(infer_output->CopyDataFromTensor(output.data()));
   EXPECT_EQ(golden, output);
+}
+
+TEST(InstanceNorm, nhwc) {
+  auto ctx = tim::vx::Context::Create();
+  auto src_graph = ctx->CreateGraph();
+
+  tim::vx::ShapeType io_shape({2, 2, 2, 2}); //nhwc
+    tim::vx::ShapeType param_shape({1});
+    tim::vx::TensorSpec input_spec(tim::vx::DataType::FLOAT32,
+                            io_shape, tim::vx::TensorAttribute::INPUT);
+    tim::vx::TensorSpec param_spec(tim::vx::DataType::FLOAT32,
+                            param_shape, tim::vx::TensorAttribute::INPUT);
+    tim::vx::TensorSpec output_spec(tim::vx::DataType::FLOAT32,
+                            io_shape, tim::vx::TensorAttribute::OUTPUT);
+
+    auto input_tensor = src_graph->CreateTensor(input_spec);
+    auto beta_tensor = src_graph->CreateTensor(param_spec);
+    auto gamma_tensor = src_graph->CreateTensor(param_spec);
+    auto output_tensor = src_graph->CreateTensor(output_spec);
+
+    std::vector<float> in_data = {
+        0.0f, 1.0f, 0.0f, 2.0f, 0.0f, 2.0f, 0.0f, 4.0f, 1.0f, -1.0f, -1.0f, 2.0f, -1.0f, -2.0f, 1.0f, 4.0f
+    };
+    std::vector<float> beta = {0};
+    std::vector<float> gamma = {1.0f};
+    std::vector<float> golden = {
+        0.0f, -1.1470304f, 0.0f, -0.22940612f, 0.0f, -0.22940612f, 0.0f, 1.6058424f, 0.99995005f,
+        -0.7337929f, -0.99995005f, 0.52413774f, -0.99995005f, -1.1531031f, 0.99995005f, 1.3627582f,
+    };
+    auto op = src_graph->CreateOperation<tim::vx::ops::InstanceNormalization>(1e-4f, tim::vx::DataLayout::CWHN);
+    (*op).BindInputs({input_tensor, beta_tensor, gamma_tensor}).BindOutputs({output_tensor});
+  // Do layout inference
+  auto transform = tim::transform::LayoutInference(src_graph, ctx);
+  auto infer_graph = transform.first;
+  auto graph_io_map = transform.second;
+  infer_graph->Compile();
+
+  auto infer_input = graph_io_map[src_graph->InputsTensor()[0]];
+  auto infer_beta = graph_io_map[src_graph->InputsTensor()[1]];
+  auto infer_gamma = graph_io_map[src_graph->InputsTensor()[2]];
+  auto infer_output = graph_io_map[src_graph->OutputsTensor()[0]];
+
+  infer_input->CopyDataToTensor(in_data.data(), in_data.size() * sizeof(float));
+  infer_beta->CopyDataToTensor(beta.data(), beta.size() * sizeof(float));
+  infer_gamma->CopyDataToTensor(gamma.data(), gamma.size() * sizeof(float));
+  infer_graph->Run();
+
+  std::vector<float> output(golden.size());
+  EXPECT_TRUE(infer_output->CopyDataFromTensor(output.data()));
+  EXPECT_TRUE(ArraysMatch(golden, output, 1e-5f));
 }

src/tim/transform/ops/instance_norm_layout_inference.h

@@ -0,0 +1,100 @@
+/****************************************************************************
+ *
+ *    Copyright (c) 2020-2023 Vivante Corporation
+ *
+ *    Permission is hereby granted, free of charge, to any person obtaining a
+ *    copy of this software and associated documentation files (the "Software"),
+ *    to deal in the Software without restriction, including without limitation
+ *    the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ *    and/or sell copies of the Software, and to permit persons to whom the
+ *    Software is furnished to do so, subject to the following conditions:
+ *
+ *    The above copyright notice and this permission notice shall be included in
+ *    all copies or substantial portions of the Software.
+ *
+ *    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ *    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ *    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ *    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ *    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ *    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ *    DEALINGS IN THE SOFTWARE.
+ *
+ *****************************************************************************/
+#ifndef TIM_LAYOUT_INFER_INSTANCE_NORM_LAYOUT_INFERENCE_H_
+#define TIM_LAYOUT_INFER_INSTANCE_NORM_LAYOUT_INFERENCE_H_
+
+#include "tim/vx/ops/instancenormalization.h"
+
+#include "ops/op_layout_inference.h"
+#include "permute_vector.h"
+#include "builtin_op_impl.h"
+
+namespace tim {
+namespace transform {
+
+class InstanceNormLayoutInfer : public OpLayoutInfer {
+ public:
+  InstanceNormLayoutInfer(
+      const std::shared_ptr<vx::Operation> op,
+      std::shared_ptr<layout_inference_impl::LayoutInferContext>& context)
+      : OpLayoutInfer(op, context) {}
+
+  // reverse any applied permute on it's input tensor
+  void OnInputs(
+      std::vector<std::shared_ptr<vx::Tensor>>& next_tensors) override {
+    vx::DataLayout layout = op_->impl()->layout_;
+    auto input_tensors = op_->impl()->InputsTensor();
+    std::shared_ptr<IPermuteVector> required_pv;
+    switch (layout)
+    { // kernel layout must be IWHO in tflite & nnapi
+      case vx::DataLayout::CWHN:
+        required_pv = std::make_shared<PermuteVector<4>>(kCWHN2WHCN);
+        break;
+      case vx::DataLayout::WHCN:
+        required_pv = MakeShared(4);
+        break;
+      default:
+        VSILOGE("The layout of input is not support.");
+        required_pv = MakeShared(4);
+        break;
+    }
+    auto input_pv = context_->GetPermuteVector(input_tensors[0]);
+    auto final_pv = input_pv->Reverse()->Add(required_pv);
+    std::shared_ptr<vx::Tensor> infer_input;
+    if (!final_pv->IsAligned()) {
+      infer_input = InsertPermute(context_->GetMapedTensor(input_tensors[0]), final_pv);
+      context_->SetPermuteVector(input_tensors[0], required_pv);
+    } else {
+      infer_input = context_->GetMapedTensor(input_tensors[0]);
+      context_->SetPermuteVector(input_tensors[0], input_pv);
+    }
+    context_->UpdateTensorMap(input_tensors[0], infer_input);
+
+    for (const auto& t_src : op_->impl()->InputsTensor()) {
+      if(t_src->IsConstTensor()) {
+        std::vector<uint8_t> dataRef(t_src->GetSpec().GetByteSize());
+        t_src->CopyDataFromTensor(dataRef.data());
+        auto t_infer = context_->infer_graph_->CreateTensor(
+                t_src->GetSpec(), (const void*)dataRef.data());
+        context_->SetPermuteVector(t_src, MakeShared(t_src->GetShape().size()));
+        context_->UpdateTensorMap(t_src, t_infer);
+      }
+    }
+
+    auto instance_norm = op_->Clone(context_->infer_graph_);
+    auto outs_infer = CreateOutputsTensor(required_pv);
+    for (const auto& i_src : op_->impl()->InputsTensor()) {
+      (*instance_norm).BindInput(context_->GetMapedTensor(i_src));
+    }
+    (*instance_norm).BindOutput(outs_infer[0]);
+    context_->SetPermuteVector(op_->impl()->OutputsTensor()[0], required_pv);
+    // Add out tensor of src_graph into next_tensor
+    next_tensors.push_back(op_->impl()->OutputsTensor()[0]);
+  }
+};
+
+}  // namespace transform
+}  // namespace tim
+
+#endif

src/tim/vx/ops/instancenormalization.cc

@@ -29,14 +29,14 @@
 namespace tim {
 namespace vx {
 namespace ops {
-InstanceNormalization::InstanceNormalization(Graph* graph, float eps)
-    : BuiltinOp(graph, VSI_NN_OP_INSTANCE_NORM), eps_(eps) {
+InstanceNormalization::InstanceNormalization(Graph* graph, float eps, DataLayout input_layout)
+    : BuiltinOp(graph, VSI_NN_OP_INSTANCE_NORM, 0, 0, input_layout), eps_(eps) {
   this->impl()->node()->nn_param.instancenorm.eps = eps_;
 }
 
 std::shared_ptr<Operation> InstanceNormalization::Clone(
     std::shared_ptr<Graph>& graph) const {
-  return graph->CreateOperation<InstanceNormalization>(this->eps_);
+  return graph->CreateOperation<InstanceNormalization>(this->eps_, this->impl_->layout_);
 }
 
 }  // namespace ops

src/tim/vx/ops/instancenormalization_test.cc

@@ -27,6 +27,48 @@
 #include "test_utils.h"
 #include "gtest/gtest.h"
 
+TEST(InstanceNorm, shape_2_2_2_2_float) {
+    auto ctx = tim::vx::Context::Create();
+    auto graph = ctx->CreateGraph();
+
+    tim::vx::ShapeType io_shape({2, 2, 2, 2}); //nchw
+    tim::vx::ShapeType param_shape({1});
+    tim::vx::TensorSpec input_spec(tim::vx::DataType::FLOAT32,
+                            io_shape, tim::vx::TensorAttribute::INPUT);
+    tim::vx::TensorSpec param_spec(tim::vx::DataType::FLOAT32,
+                            param_shape, tim::vx::TensorAttribute::INPUT);
+    tim::vx::TensorSpec output_spec(tim::vx::DataType::FLOAT32,
+                            io_shape, tim::vx::TensorAttribute::OUTPUT);
+
+    auto input_tensor = graph->CreateTensor(input_spec);
+    auto gamma_tensor = graph->CreateTensor(param_spec);
+    auto beta_tensor = graph->CreateTensor(param_spec);
+    auto output_tensor = graph->CreateTensor(output_spec);
+
+    std::vector<float> in_data = {
+        0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 2.0f, 2.0f, 4.0f, 1.0f, -1.0f, -1.0f, 1.0f, -1.0f, 2.0f, -2.0f, 4.0f
+    };
+    std::vector<float> gamma = {1.0f};
+    std::vector<float> beta = {0};
+    std::vector<float> golden = {
+        0.0f, 0.0f, 0.0f, 0.0f, -1.1470304f, -0.22940612f, -0.22940612f, 1.6058424f, 0.99995005f,
+        -0.99995005f, -0.99995005f, 0.99995005f, -0.7337929f, 0.52413774f, -1.1531031f, 1.3627582f
+    };
+    EXPECT_TRUE(input_tensor->CopyDataToTensor(in_data.data(), in_data.size() * sizeof(float)));
+    EXPECT_TRUE(gamma_tensor->CopyDataToTensor(gamma.data(), gamma.size() * sizeof(float)));
+    EXPECT_TRUE(beta_tensor->CopyDataToTensor(beta.data(), beta.size() * sizeof(float)));
+
+    auto op = graph->CreateOperation<tim::vx::ops::InstanceNormalization>(1e-4f);
+    (*op).BindInputs({input_tensor, beta_tensor, gamma_tensor}).BindOutputs({output_tensor});
+
+    EXPECT_TRUE(graph->Compile());
+    EXPECT_TRUE(graph->Run());
+
+    std::vector<float> output(16);
+    EXPECT_TRUE(output_tensor->CopyDataFromTensor(output.data()));
+    EXPECT_TRUE(ArraysMatch(golden, output, 1e-5f));
+}
+
 TEST(InstanceNorm, shape_3_6_1_float) {
     auto ctx = tim::vx::Context::Create();
     auto graph = ctx->CreateGraph();