[Tensor] Add multiple axes support

This patch adds multiple axes support for sum and average.
which is esentially needed to build bn_layer

**Self evaluation:**
1. Build test: [X]Passed [ ]Failed [ ]Skipped
2. Run test: [X]Passed [ ]Failed [ ]Skipped

Signed-off-by: Jihoon Lee <jhoon.it.lee@samsung.com>

diff --git a/nntrainer/include/tensor.h b/nntrainer/include/tensor.h
index ee6ba1e825a5871dfeef7536419bd567885c6fd6..c3129065f935317a718767dc4f5a88cec32fbd4d 100644 (file)
--- a/nntrainer/include/tensor.h
+++ b/nntrainer/include/tensor.h
@@ -306,7 +306,16 @@ public:
    * @param[in] alpha Scale the sum by this value
    * @retval    Calculated Tensor
    */
-  Tensor sum(int axis, float alpha = 1.0) const;
+  Tensor sum(unsigned int axis, float alpha = 1.0) const;
+
+  /**
+   * @brief sum all the Tensor by multiple axes
+   *
+   * @param axes axes to sum along
+   * @param alpha Scale the sum by this value
+   * @return Tensor
+   */
+  Tensor sum(const std::vector<unsigned int> &axes, float alpha = 1.0) const;
 
   /**
    * @brief     Averaging the Tensor elements according to the axis
@@ -316,7 +325,15 @@ public:
    *            3 : width direction
    * @retval    Calculated Tensor
    */
-  Tensor average(int axis) const;
+  Tensor average(unsigned int axis) const;
+
+  /**
+   * @brief average all the Tensor by multiple axes
+   *
+   * @param axes axes to sum along
+   * @return Tensor
+   */
+  Tensor average(const std::vector<unsigned int> &axes) const;
 
   /**
    * @brief     Averaging the Tensor elements by all axis

diff --git a/nntrainer/include/tensor_dim.h b/nntrainer/include/tensor_dim.h
index 4a39b13378bb1a5bc9dc1851ddd0534e8bd127db..7b8c2746b8ca913614d47c68cfd34649ee6a0f47 100644 (file)
--- a/nntrainer/include/tensor_dim.h
+++ b/nntrainer/include/tensor_dim.h
@@ -84,7 +84,7 @@ public:
   const unsigned int *getDim() const { return dim; }
   unsigned int getNumDim() const { return MAXDIM; }
 
-  unsigned int getTensorDim(unsigned int idx);
+  const unsigned int getTensorDim(unsigned int idx) const;
   void setTensorDim(unsigned int idx, unsigned int value);
   int setTensorDim(std::string input_shape);
 

diff --git a/nntrainer/src/tensor.cpp b/nntrainer/src/tensor.cpp
index ca3eff6e5fabf709b4c15230428d12f40a4c46d1..69be9b65822ac4db5919c8b707a1b43c1b4b24ca 100644 (file)
--- a/nntrainer/src/tensor.cpp
+++ b/nntrainer/src/tensor.cpp
@@ -335,13 +335,13 @@ Tensor Tensor::sum_by_batch() {
 /**
  * @brief Calculate sum according to the axis.
  */
-Tensor Tensor::sum(int axis, float alpha) const {
+Tensor Tensor::sum(unsigned int axis, float alpha) const {
   Tensor ret;
 
   const float *data = getData();
 
   if (axis >= 4)
-    throw std::out_of_range("Error: Dimension cannot exceed 3");
+    throw std::out_of_range("Error: axis is invalid");
 
   if (dim.getDim()[axis] == 1 and alpha == 1.0)
     return this->clone();
@@ -401,6 +401,18 @@ Tensor Tensor::sum(int axis, float alpha) const {
   return ret;
 }
 
+Tensor Tensor::sum(const std::vector<unsigned int> &axes, float alpha) const {
+  if (axes.empty())
+    throw std::invalid_argument("empty axes given");
+
+  Tensor ret = this->sum(axes[0], alpha);
+
+  for (unsigned int i = 1; i < axes.size(); ++i)
+    ret = ret.sum(axes[i]);
+
+  return ret;
+}
+
 /**
  * @note: This dot product flattens the fist 3 axis for the purpose of
  * computation. So, while performing, these matrices are behaving as 2-D
@@ -640,7 +652,11 @@ void Tensor::read(std::ifstream &file) {
 /**
  * @brief Calculate average value according to the axis.
  */
-Tensor Tensor::average(int axis) const {
+Tensor Tensor::average(unsigned int axis) const {
+  if (axis >= MAXDIM)
+    throw std::out_of_range(
+      "negative axis or axis more then MAXDIM is invalid");
+
   unsigned int axis_size = dim.getDim()[axis];
   if (axis_size == 1)
     return this->clone();
@@ -648,6 +664,21 @@ Tensor Tensor::average(int axis) const {
   return this->sum(axis, 1.0 / ((float)axis_size));
 }
 
+Tensor Tensor::average(const std::vector<unsigned int> &axes) const {
+  if (axes.empty())
+    return this->average();
+
+  TensorDim ret_shape;
+  for (const auto &idx : axes) {
+    if (idx >= MAXDIM) {
+      throw std::out_of_range("axis more then MAXDIM is invalid");
+    }
+    ret_shape.setTensorDim(idx, dim.getTensorDim(idx));
+  }
+
+  return this->sum(axes, 1.0 / (float)ret_shape.getDataLen());
+}
+
 /**
  * @brief Calculate average value according to the axis.
  */

diff --git a/nntrainer/src/tensor_dim.cpp b/nntrainer/src/tensor_dim.cpp
index 039116cedd8704d3ccbcbe06ece408b00ebfcb4a..fe9afc1fe709eb3b6b9fc58390a40a7515ae47e0 100644 (file)
--- a/nntrainer/src/tensor_dim.cpp
+++ b/nntrainer/src/tensor_dim.cpp
@@ -45,7 +45,7 @@ void TensorDim::resetLen() {
   len = dim[0] * feature_len;
 }
 
-unsigned int TensorDim::getTensorDim(unsigned int idx) {
+const unsigned int TensorDim::getTensorDim(unsigned int idx) const {
   if (idx > MAXDIM)
     throw std::invalid_argument(
       "[TensorDim] Tensor Dimension index should be between 0 and 4");

diff --git a/test/unittest/unittest_nntrainer_tensor.cpp b/test/unittest/unittest_nntrainer_tensor.cpp
index f920d61497ef44945ed4a47ba270c222313e75f6..532158264c66e13488cd3091f0ad646dc276f49a 100644 (file)
--- a/test/unittest/unittest_nntrainer_tensor.cpp
+++ b/test/unittest/unittest_nntrainer_tensor.cpp
@@ -676,6 +676,18 @@ TEST(nntrainer_Tensor, sum_01_n) {
   EXPECT_THROW({ input.sum(4); }, std::out_of_range);
 }
 
+TEST(nntrainer_Tensor, sum_02_n) {
+  int batch = 3;
+  int channel = 1;
+  int height = 3;
+  int width = 10;
+
+  nntrainer::Tensor input(batch, channel, height, width);
+  GEN_TEST_INPUT(input, i * (batch * height) + j * (width) + k);
+
+  EXPECT_THROW({ input.sum(-1); }, std::out_of_range);
+}
+
 TEST(nntrainer_Tensor, sum_02_p) {
   int status = ML_ERROR_NONE;
   int batch = 3;
@@ -791,6 +803,115 @@ TEST(nntrainer_Tensor, sum_03_p) {
   EXPECT_EQ(status, ML_ERROR_NONE);
 }
 
+TEST(nntrainer_Tensor, multiple_sum_invalid_args_01_n) {
+  nntrainer::Tensor t = constant(1.0, 1, 1, 1, 1);
+  EXPECT_THROW(t.sum(std::vector<unsigned int>()), std::invalid_argument);
+}
+
+TEST(nntrainer_Tensor, multiple_sum_out_of_range_n) {
+  nntrainer::Tensor t = constant(1.0, 1, 1, 1, 1);
+  EXPECT_THROW(t.sum({7}), std::out_of_range);
+}
+
+TEST(nntrainer_Tensor, multiple_sum_p) {
+  nntrainer::Tensor t = constant(1.0, 2, 3, 5, 7);
+  nntrainer::Tensor actual, expected;
+
+  actual = t.sum({0, 1});
+  expected = constant(2 * 3, 1, 1, 5, 7);
+  EXPECT_EQ(actual, expected);
+
+  actual = t.sum({1, 2, 3});
+  expected = constant(3 * 5 * 7, 2, 1, 1, 1);
+  EXPECT_EQ(actual, expected);
+
+  actual = t.sum({3, 1});
+  expected = constant(7 * 3, 2, 1, 5, 1);
+  EXPECT_EQ(actual, expected);
+
+  actual = t.sum({3, 1}, 0.5);
+  expected = constant(7 * 3 * 0.5, 2, 1, 5, 1);
+  EXPECT_EQ(actual, expected);
+}
+
+TEST(nntrainer_Tensor, average_p) {
+  nntrainer::Tensor t = constant(1.0, 2, 3, 5, 7);
+
+  nntrainer::Tensor actual, expected;
+
+  actual = t.average();
+  expected = constant(1.0, 1, 1, 1, 1);
+  EXPECT_EQ(actual, expected);
+
+  int idx = 0;
+  t = t.apply([&](float in) { return idx++ % 2; });
+
+  actual = t.average();
+  expected = constant(0.5, 1, 1, 1, 1);
+  EXPECT_EQ(actual, expected);
+}
+
+TEST(nntrainer_Tensor, average_axis_p) {
+  nntrainer::Tensor t = constant(1.0, 2, 2, 2, 2);
+  int idx = 0;
+  std::function<float(float)> f = [&](float in) { return idx++ % 2; };
+  t = t.apply(f);
+
+  nntrainer::Tensor actual, expected;
+
+  actual = t.average(0);
+  expected = constant(0, 1, 2, 2, 2).apply(f);
+  EXPECT_EQ(actual, expected);
+
+  actual = t.average(1);
+  expected = constant(0, 2, 1, 2, 2).apply(f);
+  EXPECT_EQ(actual, expected);
+
+  actual = t.average(2);
+  expected = constant(0, 2, 2, 1, 2).apply(f);
+  EXPECT_EQ(actual, expected);
+
+  actual = t.average(3);
+  expected = constant(0.5, 2, 2, 2, 1);
+  EXPECT_EQ(actual, expected);
+}
+
+TEST(nntrainer_Tensor, average_axis_out_of_range_01_n) {
+  nntrainer::Tensor t = constant(1.0, 2, 2, 2, 2);
+  EXPECT_THROW(t.average(-1), std::out_of_range);
+}
+
+TEST(nntrainer_Tensor, average_axis_out_of_range_02_n) {
+  nntrainer::Tensor t = constant(1.0, 2, 2, 2, 2);
+  EXPECT_THROW(t.average(7), std::out_of_range);
+}
+
+TEST(nntrainer_Tensor, average_multiple_axes_p) {
+  nntrainer::Tensor t = constant(1.0, 2, 3, 5, 7);
+  nntrainer::Tensor actual, expected;
+
+  actual = t.average({0, 1, 2});
+  expected = constant(1.0, 1, 1, 1, 7);
+  EXPECT_EQ(actual, expected);
+
+  actual = t.average({0, 1, 2, 3});
+  expected = constant(1.0, 1, 1, 1, 1);
+  EXPECT_EQ(actual, expected);
+
+  actual = t.average({3, 1});
+  expected = constant(1.0, 2, 1, 5, 1);
+  EXPECT_EQ(actual, expected);
+
+  actual = t.average({3, 1, 1, 1, 3});
+  expected = constant(1.0, 2, 1, 5, 1);
+  EXPECT_EQ(actual, expected);
+}
+
+TEST(nntrainer_Tensor, average_multiple_axes_01_n) {
+  nntrainer::Tensor t = constant(1.0, 2, 3, 5, 7);
+  EXPECT_THROW(t.average({5, 7}), std::out_of_range);
+}
+
 TEST(nntrainer_Tensor, dot_01_n) {
   nntrainer::Tensor input(2, 3, 4, 5);
   nntrainer::Tensor m(1, 3, 4, 5);