[tensor] Reduce num of operations for tensor sum

This patch reduces the number of operations required for tensor sum by
accumulating the consecutive axes in order to increase the precision of
the operation.

Signed-off-by: Parichay Kapoor <pk.kapoor@samsung.com>

diff --git a/nntrainer/tensor/tensor.cpp b/nntrainer/tensor/tensor.cpp
index 098b393..38062bd 100644 (file)
--- a/nntrainer/tensor/tensor.cpp
+++ b/nntrainer/tensor/tensor.cpp
@@ -214,6 +214,8 @@ bool Tensor::operator==(const Tensor &rhs) const {
   const float *rdata = rhs.getData();
 
   for (size_t i = 0; i < len; ++i) {
+    /** not checking sign change is intentional to avoid float calculation
+     * errors around 0 */
     if (std::isnan(data[i]) || std::isnan(rdata[i]) ||
         std::fabs(data[i] - rdata[i]) > epsilon)
       return false;
@@ -738,6 +740,14 @@ Tensor Tensor::sum(const std::vector<unsigned int> &axes, float alpha) const {
   return sum(axes, ret, alpha);
 }
 
+void Tensor::merge_axis(unsigned int axis1, unsigned int axis2) {
+  if (axis2 != axis1 + 1)
+    throw std::invalid_argument("Axis to be merged must be continuous.");
+
+  dim.setTensorDim(axis2, dim.getTensorDim(axis1) * dim.getTensorDim(axis2));
+  dim.setTensorDim(axis1, 1);
+}
+
 Tensor &Tensor::sum(const std::vector<unsigned int> &axes, Tensor &output,
                     float alpha) const {
   if (axes.empty())
@@ -746,12 +756,22 @@ Tensor &Tensor::sum(const std::vector<unsigned int> &axes, Tensor &output,
   if (axes.size() == 1) {
     this->sum(axes[0], output, alpha);
   } else {
-    Tensor ret = this->sum(axes[0], alpha);
+    /** club axes together */
+    Tensor new_reshaped = *this;
+    std::vector<unsigned int> new_axes = {axes[0]};
+    for (unsigned int i = 1; i < axes.size(); ++i) {
+      if (axes[i] == axes[i - 1] + 1) {
+        new_reshaped.merge_axis(axes[i - 1], axes[i]);
+        new_axes.back() = axes[i];
+      } else {
+        new_axes.push_back(axes[i]);
+      }
+    }
 
-    for (unsigned int i = 1; i < axes.size() - 1; ++i)
+    Tensor ret = new_reshaped.sum(new_axes[0]);
+    for (unsigned int i = 1; i < new_axes.size() - 1; ++i)
       ret = ret.sum(axes[i]);
-
-    ret.sum(axes.back(), output);
+    ret.sum(new_axes.back(), output, alpha);
   }
 
   return output;

diff --git a/nntrainer/tensor/tensor.h b/nntrainer/tensor/tensor.h
index b42fa16..b388eb7 100644 (file)
--- a/nntrainer/tensor/tensor.h
+++ b/nntrainer/tensor/tensor.h
@@ -1120,6 +1120,14 @@ private:
     deallocate();
     allocate();
   }
+
+  /**
+   * @brief Merge the given two axis for tensor at second axis inplace
+   *
+   * @param axis1 first axis to merge
+   * @param axis2 second axis to merge
+   */
+  void merge_axis(unsigned int axis1, unsigned int axis2);
 }; // namespace nntrainer
 
 /**