inference-engine/src/extension/ext_reorg_yolo.cpp

   1 // Copyright (C) 2018-2019 Intel Corporation
   2 // SPDX-License-Identifier: Apache-2.0
   3 //
   4
   5 #include "ext_list.hpp"
   6 #include "ext_base.hpp"
   7 #include <vector>
   8
   9 namespace InferenceEngine {
  10 namespace Extensions {
  11 namespace Cpu {
  12
  13 class ReorgYoloImpl: public ExtLayerBase {
  14 public:
  15     explicit ReorgYoloImpl(const CNNLayer* layer) {
  16         try {
  17             if (layer->insData.size() != 1 || layer->outData.empty())
  18                 THROW_IE_EXCEPTION << "Incorrect number of input/output edges!";
  19
  20             stride = layer->GetParamAsInt("stride");
  21
  22             addConfig(layer, {DataConfigurator(ConfLayout::PLN)}, {DataConfigurator(ConfLayout::PLN)});
  23         } catch (InferenceEngine::details::InferenceEngineException &ex) {
  24             errorMsg = ex.what();
  25         }
  26     }
  27
  28     StatusCode execute(std::vector<Blob::Ptr>& inputs, std::vector<Blob::Ptr>& outputs,
  29                        ResponseDesc *resp) noexcept override {
  30         const auto *src_data = inputs[0]->cbuffer().as<const float *>();
  31         auto *dst_data = outputs[0]->buffer().as<float *>();
  32
  33         int IW = (inputs[0]->getTensorDesc().getDims().size() > 3) ? inputs[0]->getTensorDesc().getDims()[3] : 1;
  34         int IH = (inputs[0]->getTensorDesc().getDims().size() > 2) ? inputs[0]->getTensorDesc().getDims()[2] : 1;
  35         int IC = (inputs[0]->getTensorDesc().getDims().size() > 1) ? inputs[0]->getTensorDesc().getDims()[1] : 1;
  36         int B = (inputs[0]->getTensorDesc().getDims().size() > 0) ? inputs[0]->getTensorDesc().getDims()[0] : 1;
  37
  38         int ic_off = IC / (stride * stride);
  39         int ih_off = IH * stride;
  40         int iw_off = IW * stride;
  41         for (int b = 0; b < B; b++) {
  42             for (int ic = 0; ic < IC; ic++) {
  43                 for (int ih = 0; ih < IH; ih++) {
  44                     for (int iw = 0; iw < IW; iw++) {
  45                         int dstIndex = b * IC * IH * IW + ic * IH * IW + ih * IW + iw;
  46
  47                         int oc = ic % ic_off;
  48                         int offset = ic / ic_off;
  49
  50                         int ow = iw * stride + offset % stride;
  51                         int oh = ih * stride + offset / stride;
  52
  53                         int srcIndex = b * ic_off * ih_off * iw_off + oc * ih_off * iw_off + oh * iw_off + ow;
  54
  55                         dst_data[dstIndex] = src_data[srcIndex];
  56                     }
  57                 }
  58             }
  59         }
  60         return OK;
  61     }
  62
  63 private:
  64     int stride;
  65 };
  66
  67 REG_FACTORY_FOR(ImplFactory<ReorgYoloImpl>, ReorgYolo);
  68
  69 }  // namespace Cpu
  70 }  // namespace Extensions
  71 }  // namespace InferenceEngine