convert mxnet squeeze expanddims, convert onnx squeeze unsqueeze

src/layer/expanddims.cpp

@@ -29,6 +29,7 @@ int ExpandDims::load_param(const ParamDict& pd)
     expand_w = pd.get(0, 0);
     expand_h = pd.get(1, 0);
     expand_c = pd.get(2, 0);
+    axes = pd.get(3, Mat());
 
     return 0;
 }
@@ -39,35 +40,80 @@ int ExpandDims::forward(const Mat& bottom_blob, Mat& top_blob, const Option& opt
     int h = bottom_blob.h;
     int dims = bottom_blob.dims;
 
+    bool _expand_w = false;
+    bool _expand_h = false;
+    bool _expand_c = false;
+
+    if (axes.empty())
+    {
+        _expand_w = expand_w;
+        _expand_h = expand_h;
+        _expand_c = expand_c;
+    }
+    else
+    {
+        const int* axes_ptr = axes;
+        for (int i=0; i<axes.w; i++)
+        {
+            int axis = axes_ptr[i];
+            if (axis < 0)
+                axis = dims + 1 + axis;// +1 for N-dim
+
+            if (dims == 1 && axis == 1)
+            {
+                _expand_h = true;
+            }
+            if (dims == 1 && axis == 2)
+            {
+                _expand_w = true;
+            }
+            if (dims == 2 && axis == 1)
+            {
+                _expand_c = true;
+            }
+            if (dims == 2 && axis == 2)
+            {
+                _expand_h = true;
+            }
+            if (dims == 2 && axis == 3)
+            {
+                _expand_w = true;
+            }
+        }
+    }
+
     top_blob = bottom_blob;
 
     if (dims == 1)
     {
-        if (expand_w)
+        if (_expand_w && _expand_h)
         {
-            if (expand_h)
-                top_blob = bottom_blob.reshape(1, 1, w, opt.blob_allocator);
-            else if (expand_c)
-                top_blob = bottom_blob.reshape(1, w, 1, opt.blob_allocator);
-            else
-                top_blob = bottom_blob.reshape(1, w, opt.blob_allocator);
+            top_blob = bottom_blob.reshape(1, w, 1, opt.blob_allocator);
         }
-        else if (expand_h)
+        else if (_expand_w)
         {
-            if (expand_c)
-                top_blob = bottom_blob.reshape(w, 1, 1, opt.blob_allocator);
-            else
-                top_blob = bottom_blob.reshape(w, 1, opt.blob_allocator);
+            top_blob = bottom_blob.reshape(1, w, opt.blob_allocator);
+        }
+        else if (_expand_h)
+        {
+            top_blob = bottom_blob.reshape(w, 1, opt.blob_allocator);
         }
     }
-    else if (dims == 2)
+
+    if (dims == 2)
     {
-        if (expand_w)
+        if (_expand_w)
+        {
             top_blob = bottom_blob.reshape(1, w, h, opt.blob_allocator);
-        else if (expand_h)
+        }
+        else if (_expand_h)
+        {
             top_blob = bottom_blob.reshape(w, 1, h, opt.blob_allocator);
-        else if (expand_c)
+        }
+        else if (_expand_c)
+        {
             top_blob = bottom_blob.reshape(w, h, 1, opt.blob_allocator);
+        }
     }
 
     if (top_blob.empty())

src/layer/expanddims.h

@@ -32,6 +32,7 @@ public:
     int expand_w;
     int expand_h;
     int expand_c;
+    Mat axes;
 };
 
 } // namespace ncnn

src/layer/squeeze.cpp

@@ -29,6 +29,7 @@ int Squeeze::load_param(const ParamDict& pd)
     squeeze_w = pd.get(0, 0);
     squeeze_h = pd.get(1, 0);
     squeeze_c = pd.get(2, 0);
+    axes = pd.get(3, Mat());
 
     return 0;
 }
@@ -40,28 +41,108 @@ int Squeeze::forward(const Mat& bottom_blob, Mat& top_blob, const Option& opt) c
     int channels = bottom_blob.c;
     int dims = bottom_blob.dims;
 
+    bool _squeeze_w = false;
+    bool _squeeze_h = false;
+    bool _squeeze_c = false;
+
+    if (axes.empty())
+    {
+        _squeeze_w = w == 1 && squeeze_w;
+        _squeeze_h = h == 1 && squeeze_h;
+        _squeeze_c = channels == 1 && squeeze_c;
+    }
+    else
+    {
+        const int* axes_ptr = axes;
+        for (int i=0; i<axes.w; i++)
+        {
+            int axis = axes_ptr[i];
+            if (axis < 0)
+                axis = dims + 1 + axis;// +1 for N-dim
+
+            if (dims == 1 && axis == 1)
+            {
+                _squeeze_w = w == 1;
+            }
+            if (dims == 2 && axis == 1)
+            {
+                _squeeze_h = h == 1;
+            }
+            if (dims == 2 && axis == 2)
+            {
+                _squeeze_w = w == 1;
+            }
+            if (dims == 3 && axis == 1)
+            {
+                _squeeze_c = channels == 1;
+            }
+            if (dims == 3 && axis == 2)
+            {
+                _squeeze_h = h == 1;
+            }
+            if (dims == 3 && axis == 3)
+            {
+                _squeeze_w = w == 1;
+            }
+        }
+    }
+
     top_blob = bottom_blob;
 
-    if (squeeze_c && dims == 3 && channels == 1)
+    if (dims == 1)
     {
-        if (squeeze_h && h == 1)
-            top_blob = bottom_blob.reshape(w, opt.blob_allocator);
-        else
-            top_blob = bottom_blob.reshape(w, h, opt.blob_allocator);
+        if (_squeeze_w)
+        {
+            top_blob = bottom_blob.reshape(1, opt.blob_allocator);
+        }
     }
-    else if (squeeze_h && dims >= 2 && h == 1)
+
+    if (dims == 2)
     {
-        if (squeeze_w && w == 1)
-            top_blob = bottom_blob.reshape(channels, opt.blob_allocator);
-        else
-            top_blob = bottom_blob.reshape(w, channels, opt.blob_allocator);
+        if (_squeeze_w && _squeeze_h)
+        {
+            top_blob = bottom_blob.reshape(1, opt.blob_allocator);
+        }
+        else if (_squeeze_w)
+        {
+            top_blob = bottom_blob.reshape(h, opt.blob_allocator);
+        }
+        else if (_squeeze_h)
+        {
+            top_blob = bottom_blob.reshape(w, opt.blob_allocator);
+        }
     }
-    else if (squeeze_w && dims >= 1 && w == 1)
+
+    if (dims == 3)
     {
-        if (squeeze_h && h == 1)
+        if (_squeeze_w && _squeeze_h && _squeeze_c)
+        {
+            top_blob = bottom_blob.reshape(1, opt.blob_allocator);
+        }
+        else if (_squeeze_w && _squeeze_h)
+        {
             top_blob = bottom_blob.reshape(channels, opt.blob_allocator);
-        else
+        }
+        else if (_squeeze_h && _squeeze_c)
+        {
+            top_blob = bottom_blob.reshape(w, opt.blob_allocator);
+        }
+        else if (_squeeze_w && _squeeze_c)
+        {
+            top_blob = bottom_blob.reshape(h, opt.blob_allocator);
+        }
+        else if (_squeeze_w)
+        {
             top_blob = bottom_blob.reshape(h, channels, opt.blob_allocator);
+        }
+        else if (_squeeze_h)
+        {
+            top_blob = bottom_blob.reshape(w, channels, opt.blob_allocator);
+        }
+        else if (_squeeze_c)
+        {
+            top_blob = bottom_blob.reshape(w, h, opt.blob_allocator);
+        }
     }
 
     if (top_blob.empty())

src/layer/squeeze.h

@@ -32,6 +32,7 @@ public:
     int squeeze_w;
     int squeeze_h;
     int squeeze_c;
+    Mat axes;
 };
 
 } // namespace ncnn

tools/mxnet/mxnet2ncnn.cpp

@@ -1435,6 +1435,10 @@ int main(int argc, char** argv)
         {
             fprintf(pp, "%-16s", "UnaryOp");
         }
+        else if (n.op == "squeeze")
+        {
+            fprintf(pp, "%-16s", "Squeeze");
+        }
         else if (n.op == "tan")
         {
             fprintf(pp, "%-16s", "UnaryOp");
@@ -2049,20 +2053,7 @@ int main(int argc, char** argv)
         {
             int axis = n.attr("axis");
 
-            int expand_w = 0;
-            int expand_h = 0;
-            int expand_c = 0;
-
-            if (axis == 0)
-                expand_c = 1;
-            if (axis == 1)
-                expand_h = 1;
-            if (axis == 2)
-                expand_w = 1;
-
-            fprintf(pp, " 0=%d", expand_w);
-            fprintf(pp, " 1=%d", expand_h);
-            fprintf(pp, " 2=%d", expand_c);
+            fprintf(pp, " -23303=1,%d", axis);
         }
         else if (n.op == "Flatten")
         {
@@ -2463,6 +2454,25 @@ int main(int argc, char** argv)
             int op_type = 4;
             fprintf(pp, " 0=%d", op_type);
         }
+        else if (n.op == "squeeze")
+        {
+            std::vector<int> axis = n.attr("axis");
+
+            if (axis.empty())
+            {
+                fprintf(pp, " 0=1");
+                fprintf(pp, " 1=1");
+                fprintf(pp, " 2=1");
+            }
+            else
+            {
+                fprintf(pp, " -23303=%d", axis.size());
+                for (int i=0; i<(int)axis.size(); i++)
+                {
+                    fprintf(pp, ",%d", axis[i]);
+                }
+            }
+        }
         else if (n.op == "tan")
         {
             int op_type = 11;

tools/ncnnoptimize.cpp

@@ -36,6 +36,7 @@
 #include "layer/eltwise.h"
 #include "layer/elu.h"
 #include "layer/exp.h"
+#include "layer/expanddims.h"
 #include "layer/flatten.h"
 #include "layer/hardsigmoid.h"
 #include "layer/hardswish.h"
@@ -68,6 +69,7 @@
 #include "layer/slice.h"
 #include "layer/shufflechannel.h"
 #include "layer/softmax.h"
+#include "layer/squeeze.h"
 #include "layer/threshold.h"
 #include "layer/unaryop.h"
 #include "layer/yolodetectionoutput.h"
@@ -2014,6 +2016,16 @@ int NetOptimize::save(const char* parampath, const char* binpath)
             fprintf_param_value(" 1=%e", scale)
             fprintf_param_value(" 2=%e", shift)
         }
+        else if (layer->type == "ExpandDims")
+        {
+            ncnn::ExpandDims* op = (ncnn::ExpandDims*)layer;
+            ncnn::ExpandDims* op_default = (ncnn::ExpandDims*)layer_default;
+
+            fprintf_param_value(" 0=%d", expand_w)
+            fprintf_param_value(" 1=%d", expand_h)
+            fprintf_param_value(" 2=%d", expand_c)
+            { if (!op->axes.empty()) fprintf_param_int_array(0, op->axes, pp); }
+        }
         else if (layer->type == "HardSigmoid")
         {
             ncnn::HardSigmoid* op = (ncnn::HardSigmoid*)layer;
@@ -2333,6 +2345,16 @@ int NetOptimize::save(const char* parampath, const char* binpath)
                 fprintf(pp, " 1=%d", fixbug0);
             }
         }
+        else if (layer->type == "Squeeze")
+        {
+            ncnn::Squeeze* op = (ncnn::Squeeze*)layer;
+            ncnn::Squeeze* op_default = (ncnn::Squeeze*)layer_default;
+
+            fprintf_param_value(" 0=%d", squeeze_w)
+            fprintf_param_value(" 1=%d", squeeze_h)
+            fprintf_param_value(" 2=%d", squeeze_c)
+            { if (!op->axes.empty()) fprintf_param_int_array(0, op->axes, pp); }
+        }
         else if (layer->type == "Threshold")
         {
             ncnn::Threshold* op = (ncnn::Threshold*)layer;

tools/onnx/onnx2ncnn.cpp

@@ -1227,6 +1227,10 @@ int main(int argc, char** argv)
         {
             fprintf(pp, "%-16s", "UnaryOp");
         }
+        else if (op == "Squeeze")
+        {
+            fprintf(pp, "%-16s", "Squeeze");
+        }
         else if (op == "Sub")
         {
             fprintf(pp, "%-16s", "BinaryOp");
@@ -1251,6 +1255,10 @@ int main(int argc, char** argv)
         {
             fprintf(pp, "%-16s", "Interp");
         }
+        else if (op == "Unsqueeze")
+        {
+            fprintf(pp, "%-16s", "ExpandDims");
+        }
         else
         {
             // TODO
@@ -2078,6 +2086,27 @@ int main(int argc, char** argv)
             int op_type = 5;
             fprintf(pp, " 0=%d", op_type);
         }
+        else if (op == "Squeeze")
+        {
+            std::vector<int> axes = get_node_attr_ai(node, "axes");
+
+            if (axes.empty())
+            {
+                fprintf(pp, " 0=1");
+                fprintf(pp, " 1=1");
+                fprintf(pp, " 2=1");
+            }
+            else
+            {
+                fprintf(pp, " -23303=%d", axes.size());
+                for (int i=0; i<(int)axes.size(); i++)
+                {
+                    if (axes[i] == 0 || axes[i] > 3 || axes[i] < -3)
+                        fprintf(stderr, "Unsupported squeeze axes !\n");
+                    fprintf(pp, ",%d", axes[i]);
+                }
+            }
+        }
         else if (op == "Sub")
         {
             int op_type = 1;
@@ -2201,6 +2230,18 @@ int main(int argc, char** argv)
             fprintf(pp, " 1=%e", h_scale);
             fprintf(pp, " 2=%e", w_scale);
         }
+        else if (op == "Unsqueeze")
+        {
+            std::vector<int> axes = get_node_attr_ai(node, "axes");
+
+            fprintf(pp, " -23303=%d", axes.size());
+            for (int i=0; i<(int)axes.size(); i++)
+            {
+                if (axes[i] == 0 || axes[i] > 4 || axes[i] < -4)
+                    fprintf(stderr, "Unsupported unsqueeze axes !\n");
+                fprintf(pp, ",%d", axes[i]);
+            }
+        }
         else
         {
             // TODO op specific param