fix pnnx softmax/normalize/slice negative axis conversion to ncnn (#4284

)
Tencent · Oct 19, 2022 · 777e4ef · 777e4ef
1 parent 549152c
commit 777e4ef
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 24 deletions.
diff --git a/tools/pnnx/src/pass_ncnn/F_normalize.cpp b/tools/pnnx/src/pass_ncnn/F_normalize.cpp
@@ -45,11 +45,6 @@ pnnx.Output             output      1 0 out
     {
         const int batch_index = op->inputs[0]->params["__batch_index"].i;
 
-        int input_rank = op->inputs[0]->shape.size();
-
-        if (batch_index >= 0 && batch_index < input_rank)
-            input_rank -= 1;
-
         int axis = captured_params.at("dim").i;
         if (axis == batch_index)
         {
@@ -58,7 +53,10 @@ pnnx.Output             output      1 0 out
         }
 
         if (axis < 0)
+        {
+            int input_rank = op->inputs[0]->shape.size();
             axis = input_rank + axis;
+        }
 
         if (axis > batch_index)
             axis -= 1;
@@ -75,6 +73,11 @@ pnnx.Output             output      1 0 out
             return;
         }
 
+        int input_rank = op->inputs[0]->shape.size();
+
+        if (batch_index >= 0 && batch_index < input_rank)
+            input_rank -= 1;
+
         if (input_rank == 2 || axis != 0)
         {
             fprintf(stderr, "unsupported normalize for %d-rank tensor with axis %d\n", input_rank, axis);

diff --git a/tools/pnnx/src/pass_ncnn/F_softmax.cpp b/tools/pnnx/src/pass_ncnn/F_softmax.cpp
@@ -45,11 +45,6 @@ pnnx.Output             output      1 0 out
     {
         const int batch_index = op->inputs[0]->params["__batch_index"].i;
 
-        int input_rank = op->inputs[0]->shape.size();
-
-        if (batch_index >= 0 && batch_index < input_rank)
-            input_rank -= 1;
-
         int axis = captured_params.at("dim").i;
         if (axis == batch_index)
         {
@@ -58,7 +53,10 @@ pnnx.Output             output      1 0 out
         }
 
         if (axis < 0)
+        {
+            int input_rank = op->inputs[0]->shape.size();
             axis = input_rank + axis;
+        }
 
         if (axis > batch_index)
             axis -= 1;

diff --git a/tools/pnnx/src/pass_ncnn/Tensor_slice.cpp b/tools/pnnx/src/pass_ncnn/Tensor_slice.cpp
@@ -60,15 +60,17 @@ pnnx.Output             output      1 0 out
 
         const int batch_index = op->inputs[0]->params["__batch_index"].i;
 
-        int input_rank = op->inputs[0]->shape.size();
+        {
+            int input_rank = op->inputs[0]->shape.size();
 
-        if (batch_index >= 0 && batch_index < input_rank)
-            input_rank -= 1;
+            if (batch_index >= 0 && batch_index < input_rank)
+                input_rank -= 1;
 
-        if (input_rank > 4)
-        {
-            fprintf(stderr, "slice %d-rank tensor with %d-rank axes is not possible!\n", input_rank, axes_rank);
-            return;
+            if (input_rank > 4)
+            {
+                fprintf(stderr, "slice %d-rank tensor with %d-rank axes is not possible!\n", input_rank, axes_rank);
+                return;
+            }
         }
 
         for (int i = 0; i < axes_rank; i++)
@@ -80,7 +82,10 @@ pnnx.Output             output      1 0 out
             }
 
             if (axes[i] < 0)
+            {
+                int input_rank = op->inputs[0]->shape.size();
                 axes[i] = input_rank + axes[i];
+            }
 
             if (axes[i] > batch_index)
                 axes[i] -= 1;

diff --git a/tools/pnnx/src/pass_ncnn/nn_Softmax.cpp b/tools/pnnx/src/pass_ncnn/nn_Softmax.cpp
@@ -45,11 +45,6 @@ pnnx.Output             output      1 0 out
     {
         const int batch_index = op->inputs[0]->params["__batch_index"].i;
 
-        int input_rank = op->inputs[0]->shape.size();
-
-        if (batch_index >= 0 && batch_index < input_rank)
-            input_rank -= 1;
-
         int axis = captured_params.at("dim").i;
         if (axis == batch_index)
         {
@@ -58,7 +53,10 @@ pnnx.Output             output      1 0 out
         }
 
         if (axis < 0)
+        {
+            int input_rank = op->inputs[0]->shape.size();
             axis = input_rank + axis;
+        }
 
         if (axis > batch_index)
             axis -= 1;

diff --git a/tools/pnnx/tests/ncnn/test_F_softmax.py b/tools/pnnx/tests/ncnn/test_F_softmax.py
@@ -24,7 +24,8 @@ def forward(self, x, y, z):
         x = F.softmax(x, 0)
         y = F.softmax(y, 1)
         z = F.softmax(z, 2)
-        return x, y, z
+        z2 = F.softmax(z, -1)
+        return x, y, z, z2
 
 def test():
     net = Model()

diff --git a/tools/pnnx/tests/ncnn/test_nn_Softmax.py b/tools/pnnx/tests/ncnn/test_nn_Softmax.py
@@ -23,12 +23,14 @@ def __init__(self):
         self.act_0 = nn.Softmax(dim=0)
         self.act_1 = nn.Softmax(dim=1)
         self.act_2 = nn.Softmax(dim=2)
+        self.act_3 = nn.Softmax(dim=-1)
 
     def forward(self, x, y, z):
         x = self.act_0(x)
         y = self.act_1(y)
         z = self.act_2(z)
-        return x, y, z
+        z2 = self.act_3(z)
+        return x, y, z, z2
 
 def test():
     net = Model()