Merge branch 'master' into feature/non_uniform

kendryte · Nov 15, 2023 · 351e795 · 351e795
2 parents d1f4078 + 734cce0
commit 351e795
Showing 1 changed file with 19 additions and 6 deletions.
diff --git a/src/Nncase.Passes/Rules/Neutral/AddPreProcess.cs b/src/Nncase.Passes/Rules/Neutral/AddPreProcess.cs
@@ -175,13 +175,26 @@ protected override Task<IRModule> RunCoreAsync(IRModule module, RunPassContext o
             // Normalization
             if (mean.Length != 0)
             {
-                newInput = mean.Length switch
+                Expr meanCall;
+                Expr stdCall;
+                switch (mean.Length)
                 {
-                    3 when inputShape.Length == 4 => (newInput - Tensor.From(mean, new[] { 1, mean.Length, 1, 1 })) /
-                                                     Tensor.From(std, new[] { 1, std.Length, 1, 1 }),
-                    _ => (newInput - Tensor.From(new float[] { mean[0] }, new[] { 1 })) /
-                         Tensor.From(new float[] { std[0] }, new[] { 1 }),
-                };
+                    case 3 when inputShape.Length == 4:
+                        meanCall = (Expr)Tensor.From(mean, new[] { 1, mean.Length, 1, 1 });
+                        stdCall = (Expr)Tensor.From(std, new[] { 1, std.Length, 1, 1 });
+                        break;
+
+                    default:
+                        meanCall = (Expr)Tensor.From(new float[] { mean[0] }, new[] { 1 });
+                        stdCall = (Expr)Tensor.From(new float[] { std[0] }, new[] { 1 });
+                        break;
+                }
+
+                meanCall.Metadata.OutputNames = new[] { "Mean" };
+                stdCall.Metadata.OutputNames = new[] { "Std" };
+                var subMean = (newInput - meanCall).With(metadata: new IRMetadata() { OutputNames = new[] { input.Metadata.OutputNames?[0] + "_SubMean" } });
+                var divStd = (subMean / stdCall).With(metadata: new IRMetadata() { OutputNames = new[] { input.Metadata.OutputNames?[0] + "_DivStd" } });
+                newInput = divStd;
 
                 // newInput = Binary(BinaryOp.Div, Binary(BinaryOp.Sub, newInput, Tensor.From(mean, new []{1,3,1,1})), Const.FromTensor(std) );
             }