[softmax] Bug fix for softmax

Softmax should apply on 1 axis (last axis in this patch) than all the axis.

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

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

diff --git a/nntrainer/include/tensor.h b/nntrainer/include/tensor.h
index fd0291e..3e2ce24 100644 (file)
--- a/nntrainer/include/tensor.h
+++ b/nntrainer/include/tensor.h
@@ -61,7 +61,7 @@ public:
    * @brief     Constructor of Tensor
    * @param[in] batch Batch of Tensor
    * @param[in] channel Channel of Tensor
-   * @param[in] heihgt Height of Tensor
+   * @param[in] height Height of Tensor
    * @param[in] width Width of Tensor
    */
   Tensor(int batch, int channel, int height, int width) :
@@ -70,20 +70,26 @@ public:
   /**
    * @brief     Constructor of Tensor
    * @param[in] channel Channel of Tensor
-   * @param[in] heihgt Height of Tensor
+   * @param[in] height Height of Tensor
    * @param[in] width Width of Tensor
    */
   Tensor(int channel, int height, int width) :
     Tensor(1, channel, height, width){};
 
   /**
-   * @brief     Constructor of Tensor with batch size one
-   * @param[in] heihgt Height of Tensor
+   * @brief     Constructor of Tensor with batch size one and channel size one
+   * @param[in] height Height of Tensor
    * @param[in] width Width of Tensor
    */
   Tensor(int height, int width) : Tensor(1, 1, height, width){};
 
   /**
+   * @brief     Constructor of Tensor with just width
+   * @param[in] width Width of Tensor
+   */
+  Tensor(int width) : Tensor(1, 1, 1, width){};
+
+  /**
    * @brief     Constructor of Tensor
    * @param[in] d data for the Tensor
    */

diff --git a/nntrainer/src/activation_layer.cpp b/nntrainer/src/activation_layer.cpp
index f7071ed..0486877 100644 (file)
--- a/nntrainer/src/activation_layer.cpp
+++ b/nntrainer/src/activation_layer.cpp
@@ -161,7 +161,7 @@ Tensor ActivationLayer::softmax(Tensor const &t) {
   const float *tp;
 
   Tensor result(t.getDim());
-  Tensor divisor(t.getDim());
+  Tensor divisor(t.getWidth());
 
   dp = divisor.getData();
   rp = result.getData();
@@ -170,35 +170,27 @@ Tensor ActivationLayer::softmax(Tensor const &t) {
   divisor.setZero();
 
   for (int k = 0; k < batch; k++) {
-    int index = k * channel * height * width;
-    float m = std::numeric_limits<float>::lowest();
-    // find max
     for (int c = 0; c < channel; c++) {
       for (int i = 0; i < height; i++) {
+        int index =
+          k * channel * height * width + c * height * width + i * width;
+        float m = std::numeric_limits<float>::lowest();
+
+        // find max
         for (int j = 0; j < width; j++) {
-          if (tp[index + c * height * width + i * width + j] > m)
-            m = tp[index + c * height * width + i * width + j];
+          if (tp[index + j] > m)
+            m = tp[index + j];
         }
-      }
-    }
 
-    // shiftx
-    float sum = 0.0f;
-    for (int c = 0; c < channel; c++) {
-      for (int i = 0; i < height; i++) {
+        // shiftx
+        float sum = 0.0f;
         for (int j = 0; j < width; j++) {
-          dp[index + c * height * width + width * i + j] =
-            exp(tp[index + c * height * width + i * width + j] - m);
-          sum += dp[index + c * height * width + width * i + j];
+          dp[j] = exp(tp[index + j] - m);
+          sum += dp[j];
         }
-      }
-    }
 
-    for (int c = 0; c < channel; c++) {
-      for (int i = 0; i < height; i++) {
         for (int j = 0; j < width; j++) {
-          rp[index + c * height * width + width * i + j] =
-            dp[index + c * height * width + width * i + j] / sum;
+          rp[index + j] = dp[j] / sum;
         }
       }
     }