Merge pull request #20483 from thezane:support-cumsum-layer-for-onnx

thezane · web-flow · commit 210bfaf8d68b · 2021-08-17T20:09:25.000+03:00
* Support cumsum layer for onnx

* Add unit tests

* Address review comments
diff --git a/modules/dnn/include/opencv2/dnn/all_layers.hpp b/modules/dnn/include/opencv2/dnn/all_layers.hpp
@@ -723,6 +723,15 @@ CV__DNN_INLINE_NS_BEGIN
         static Ptr<Layer> create(const LayerParams& params);
     };
 
+    class CV_EXPORTS CumSumLayer : public Layer
+    {
+    public:
+        int exclusive;
+        int reverse;
+
+        static Ptr<CumSumLayer> create(const LayerParams& params);
+    };
+
 //! @}
 //! @}
 CV__DNN_INLINE_NS_END
diff --git a/modules/dnn/src/init.cpp b/modules/dnn/src/init.cpp
@@ -140,6 +140,7 @@ void initializeLayerFactory()
 
     CV_DNN_REGISTER_LAYER_CLASS(LSTM,           LSTMLayer);
     CV_DNN_REGISTER_LAYER_CLASS(GRU,            GRULayer);
+    CV_DNN_REGISTER_LAYER_CLASS(CumSum,         CumSumLayer);
 }
 
 CV__DNN_INLINE_NS_END
diff --git a/modules/dnn/src/layers/cumsum_layer.cpp b/modules/dnn/src/layers/cumsum_layer.cpp
@@ -0,0 +1,131 @@
+// This file is part of OpenCV project.
+// It is subject to the license terms in the LICENSE file found in the top-level directory
+// of this distribution and at https://opencv.org/license.html.
+
+#include "../precomp.hpp"
+#include "layers_common.hpp"
+
+#include <opencv2/dnn/shape_utils.hpp>
+
+namespace cv
+{
+namespace dnn
+{
+
+class CumSumLayerImpl CV_FINAL : public CumSumLayer
+{
+public:
+    CumSumLayerImpl(const LayerParams &params)
+    {
+        axis_raw = params.get<int>("axis", 0);
+        exclusive_raw = params.get<int>("exclusive", 0);
+        reverse_raw = params.get<int>("reverse", 0);
+        setParamsFrom(params);
+    }
+
+    bool getMemoryShapes(const std::vector<MatShape> &inputs,
+                         const int requiredOutputs,
+                         std::vector<MatShape> &outputs,
+                         std::vector<MatShape> &internals) const CV_OVERRIDE
+    {
+        Layer::getMemoryShapes(inputs, requiredOutputs, outputs, internals);
+        return true;
+    }
+
+    void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr, OutputArrayOfArrays internals_arr) CV_OVERRIDE
+    {
+        CV_TRACE_FUNCTION();
+        CV_TRACE_ARG_VALUE(name, "name", name.c_str());
+
+        if (inputs_arr.depth() == CV_16S)
+        {
+            forward_fallback(inputs_arr, outputs_arr, internals_arr);
+            return;
+        }
+
+        std::vector<Mat> inputs, outputs, internals;
+        inputs_arr.getMatVector(inputs);
+        outputs_arr.getMatVector(outputs);
+
+        // Get x tensor.
+        const auto &src_mat = inputs[0];
+        const auto *src_ptr = src_mat.ptr<float>();
+
+        // Get axis.
+        const int axis = normalize_axis(axis_raw, src_mat.dims);
+
+        // Get y tensor.
+        auto &dst_mat = outputs[0];
+        src_mat.copyTo(dst_mat);
+        auto *dst_ptr = dst_mat.ptr<float>();
+
+        // Get flags.
+        const auto exclusive = exclusive_raw == 1;
+        const auto reverse = reverse_raw == 1;
+
+        // Get parameters to iterate outer dimension.
+        const size_t outer_size = src_mat.total(0, axis);
+        const size_t outer_step_length = src_mat.total(axis);
+
+        // Get parameters to iterate inner dimension.
+        const size_t inner_size = src_mat.size[axis];
+
+        if (!inner_size)
+            return;
+
+        const size_t inner_step_length = src_mat.total(axis + 1);
+        const int inner_step = (reverse ? -1 : 1) * inner_step_length;
+        const int inner_start = reverse ? inner_size - 1 : 0;
+        const int inner_stop = reverse ? -1 : inner_size;
+        const int inner_delta = reverse ? -1 : 1;
+
+        // Get parameters to populate channels.
+        const size_t num_channels = src_mat.total(axis + 1);
+
+        for (size_t outer_dim = 0; outer_dim < outer_size; outer_dim++)
+        {
+            const size_t outer_offset = outer_dim * outer_step_length;
+            size_t src_offset = outer_offset + inner_start * inner_step_length;
+
+            // Populate first element of inner dimension.
+            for (size_t channel = 0; channel < num_channels; channel++)
+            {
+                if (exclusive)
+                {
+                    dst_ptr[src_offset + channel] = 0.0f;
+                }
+                else
+                {
+                    dst_ptr[src_offset + channel] = src_ptr[src_offset + channel];
+                    src_offset += inner_step;
+                }
+            }
+
+            // Populate remaining elements of inner dimension.
+            for (int inner_dim = inner_start + inner_delta; inner_dim != inner_stop; inner_dim += inner_delta)
+            {
+                const size_t dst_offset = outer_offset + inner_dim * inner_step_length;
+
+                for (size_t channel = 0; channel < num_channels; channel++)
+                {
+                    const size_t previous_dst_offset = dst_offset - inner_step;
+                    dst_ptr[dst_offset + channel] = dst_ptr[previous_dst_offset + channel] +
+                            src_ptr[src_offset + channel];
+                    src_offset += inner_step;
+                }
+            }
+        }
+    }
+
+    int axis_raw;
+    int exclusive_raw;
+    int reverse_raw;
+};
+
+Ptr<CumSumLayer> CumSumLayer::create(const LayerParams& params)
+{
+    return Ptr<CumSumLayer>(new CumSumLayerImpl(params));
+}
+
+}
+}
diff --git a/modules/dnn/src/onnx/onnx_importer.cpp b/modules/dnn/src/onnx/onnx_importer.cpp
@@ -177,6 +177,7 @@ class ONNXImporter
     void parseUpsample             (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseSoftMax              (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseDetectionOutput      (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
+    void parseCumSum               (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseCustom               (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
 
     const DispatchMap dispatch;
@@ -641,7 +642,8 @@ const std::set<String>& ONNXImporter::getSupportedTypes()
         "Dropout",
         "Identity",
         "Crop",
-        "Normalize"
+        "Normalize",
+        "CumSum"
     };
     return layerTypes;
 }
@@ -2399,6 +2401,23 @@ void ONNXImporter::parseDetectionOutput(LayerParams& layerParams, const opencv_o
     addLayer(layerParams, node_proto);
 }
 
+void ONNXImporter::parseCumSum(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
+{
+    layerParams.type = "CumSum";
+
+    // Get axis.
+    const std::string& input1 = node_proto.input(1);
+
+    if (constBlobs.find(input1) != constBlobs.end())
+    {
+        Mat axis_blob = getBlob(input1);
+        CV_Assert(axis_blob.total() == 1u);
+        layerParams.set("axis", axis_blob.at<int>(0));
+    }
+
+    addLayer(layerParams, node_proto);
+}
+
 void ONNXImporter::parseCustom(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
 {
     for (int j = 0; j < node_proto.input_size(); j++) {
@@ -2456,6 +2475,7 @@ const ONNXImporter::DispatchMap ONNXImporter::buildDispatchMap()
     dispatch["Upsample"] = &ONNXImporter::parseUpsample;
     dispatch["SoftMax"] = dispatch["LogSoftmax"] = &ONNXImporter::parseSoftMax;
     dispatch["DetectionOutput"] = &ONNXImporter::parseDetectionOutput;
+    dispatch["CumSum"] = &ONNXImporter::parseCumSum;
     dispatch["Custom"] = &ONNXImporter::parseCustom;
 
     return dispatch;
diff --git a/modules/dnn/test/test_onnx_importer.cpp b/modules/dnn/test/test_onnx_importer.cpp
@@ -1362,6 +1362,15 @@ TEST_P(Test_ONNX_nets, Resnet34_kinetics)
     expectNoFallbacksFromIE(net);
 }
 
+TEST_P(Test_ONNX_layers, CumSum)
+{
+    testONNXModels("cumsum_1d_exclusive_1");
+    testONNXModels("cumsum_1d_reverse");
+    testONNXModels("cumsum_1d_exclusive_1_reverse");
+    testONNXModels("cumsum_2d_dim_1");
+    testONNXModels("cumsum_3d_dim_2");
+}
+
 INSTANTIATE_TEST_CASE_P(/**/, Test_ONNX_nets, dnnBackendsAndTargets());
 
 }} // namespace

Original file line number	Diff line number	Diff line change
`@@ -140,6 +140,7 @@ void initializeLayerFactory()`
`140`	`140`
`141`	`141`	`CV_DNN_REGISTER_LAYER_CLASS(LSTM, LSTMLayer);`
`142`	`142`	`CV_DNN_REGISTER_LAYER_CLASS(GRU, GRULayer);`
	`143`	`+ CV_DNN_REGISTER_LAYER_CLASS(CumSum, CumSumLayer);`
`143`	`144`	`}`
`144`	`145`
`145`	`146`	`CV__DNN_INLINE_NS_END`