modules/dnn/src/cuda4dnn/primitives/reshape.hpp

   1 // This file is part of OpenCV project.
   2 // It is subject to the license terms in the LICENSE file found in the top-level directory
   3 // of this distribution and at http://opencv.org/license.html.
   4
   5 #ifndef OPENCV_DNN_SRC_CUDA4DNN_PRIMITIVES_RESHAPE_HPP
   6 #define OPENCV_DNN_SRC_CUDA4DNN_PRIMITIVES_RESHAPE_HPP
   7
   8 #include "../../op_cuda.hpp"
   9
  10 #include "../csl/stream.hpp"
  11 #include "../csl/tensor.hpp"
  12 #include "../csl/tensor_ops.hpp"
  13
  14 #include <utility>
  15
  16 namespace cv { namespace dnn { namespace cuda4dnn {
  17
  18     template <class T>
  19     class ReshapeOp final : public CUDABackendNode {
  20     public:
  21         using wrapper_type = GetCUDABackendWrapperType<T>;
  22
  23         ReshapeOp(csl::Stream stream_) : stream(std::move(stream_)) { }
  24
  25         void forward(
  26             const std::vector<cv::Ptr<BackendWrapper>>& inputs,
  27             const std::vector<cv::Ptr<BackendWrapper>>& outputs,
  28             csl::Workspace& workspace) override
  29         {
  30             /* sometimes the output shape is passed as extra inputs; hence, >= instead of == */
  31             CV_Assert(inputs.size() >= outputs.size());
  32
  33             for (int i = 0; i < outputs.size(); i++)
  34             {
  35                 auto input_wrapper = inputs[i].dynamicCast<wrapper_type>();
  36                 auto input = input_wrapper->getView();
  37
  38                 auto output_wrapper = outputs[i].dynamicCast<wrapper_type>();
  39                 auto output = output_wrapper->getSpan();
  40
  41                 if (input.get() != output.get())
  42                 {
  43                     while (input.rank() < output.rank())
  44                         input.unsqueeze();
  45
  46                     while (output.rank() < input.rank())
  47                         output.unsqueeze();
  48
  49                     input.reshape_as(output);
  50                     csl::tensor_ops::copy(stream, output, input);
  51                 }
  52             }
  53         }
  54
  55     private:
  56         csl::Stream stream;
  57     };
  58
  59 }}} /* namespace cv::dnn::cuda4dnn */
  60
  61 #endif /* OPENCV_DNN_SRC_CUDA4DNN_PRIMITIVES_RESHAPE_HPP */