fix LN evaluator (#1131)

kendryte · Nov 21, 2023 · f6e4674 · f6e4674
1 parent 2571d16
commit f6e4674
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/src/Nncase.Evaluator/NN/LayerNorm.cs b/src/Nncase.Evaluator/NN/LayerNorm.cs
@@ -145,7 +145,7 @@ private UInt128 GetRingReduceCommunicate(DistributedType distributedType, int[]
 #if true
     private Tensor LayerNormImpl(Tensor input, Tensor scale, Tensor bias, int axis, float epsilon, bool useMean = true)
     {
-        int outputSize = 1;
+        int outerSize = 1;
         int innerSize = 1;
         float[] inputArray = input.ToArray<float>();
         float[] outputArray = new float[inputArray.Length];
@@ -157,23 +157,25 @@ private Tensor LayerNormImpl(Tensor input, Tensor scale, Tensor bias, int axis,
 
         for (int i = 0; i < axis; i++)
         {
-            outputSize *= inShape[i];
+            outerSize *= inShape[i];
         }
 
         for (int i = axis; i < inShape.Length; i++)
         {
             innerSize *= inShape[i];
         }
 
-        for (int batch = 0; batch < outputSize; batch++)
+        for (int batch = 0; batch < outerSize; batch++)
         {
             float mean1 = 0f;
             if (useMean)
             {
                 for (int i = 0; i < innerSize; i++)
                 {
-                    mean1 += inputArray[(i + (batch * innerSize)) % inputArray.Length] / innerSize;
+                    mean1 += inputArray[(i + (batch * innerSize)) % inputArray.Length];
                 }
+
+                mean1 /= innerSize;
             }
 
             float[] sub = new float[innerSize];
@@ -191,9 +193,11 @@ private Tensor LayerNormImpl(Tensor input, Tensor scale, Tensor bias, int axis,
             float mean2 = 0f;
             for (int i = 0; i < innerSize; i++)
             {
-                mean2 += pow[i] / innerSize;
+                mean2 += pow[i];
             }
 
+            mean2 /= innerSize;
+
             float add = mean2 + epsilon;
             float sqrt = (float)System.Math.Sqrt(add);