diff --git a/src/layer/expanddims.cpp b/src/layer/expanddims.cpp index 9daecda6a..37c0026a9 100644 --- a/src/layer/expanddims.cpp +++ b/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= 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()) diff --git a/src/layer/squeeze.h b/src/layer/squeeze.h index cf8a2a164..cea5a413c 100644 --- a/src/layer/squeeze.h +++ b/src/layer/squeeze.h @@ -32,6 +32,7 @@ public: int squeeze_w; int squeeze_h; int squeeze_c; + Mat axes; }; } // namespace ncnn diff --git a/tools/mxnet/mxnet2ncnn.cpp b/tools/mxnet/mxnet2ncnn.cpp index e2e08ece6..76d69a677 100644 --- a/tools/mxnet/mxnet2ncnn.cpp +++ b/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 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; diff --git a/tools/ncnnoptimize.cpp b/tools/ncnnoptimize.cpp index 588a14e79..7829fd485 100644 --- a/tools/ncnnoptimize.cpp +++ b/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; diff --git a/tools/onnx/onnx2ncnn.cpp b/tools/onnx/onnx2ncnn.cpp index 97c82933a..f4b6ab737 100644 --- a/tools/onnx/onnx2ncnn.cpp +++ b/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 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 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