inference-engine/src/extension/ext_unsqueeze.cpp

   1 // Copyright (C) 2019 Intel Corporation
   2 // SPDX-License-Identifier: Apache-2.0
   3 //
   4
   5 #include "ext_list.hpp"
   6 #include "ext_base.hpp"
   7
   8 #include <cmath>
   9 #include <string>
  10 #include <vector>
  11 #include <cassert>
  12 #include "ie_parallel.hpp"
  13
  14 namespace InferenceEngine {
  15 namespace Extensions {
  16 namespace Cpu {
  17
  18 class UnsqueezeImpl: public ExtLayerBase {
  19 public:
  20     explicit UnsqueezeImpl(const CNNLayer* layer) {
  21         try {
  22             if (layer->insData.empty() || layer->outData.empty())
  23                 THROW_IE_EXCEPTION << layer->name << " Incorrect number of input/output edges!";
  24
  25             if (layer->insData.size() != 2)
  26                 THROW_IE_EXCEPTION << layer->name << " Incorrect number of input edges!";
  27
  28             idx_dims = layer->insData[UNSQUEEZE_INDEXES].lock()->getTensorDesc().getDims();
  29             data_dims = layer->insData[UNSQUEEZE_DATA].lock()->getTensorDesc().getDims();
  30             if (idx_dims.size() > 1)
  31                 THROW_IE_EXCEPTION << layer->name << " Index vector should be 1 dimension";
  32
  33             if (layer->insData[UNSQUEEZE_INDEXES].lock()->getTensorDesc().getPrecision() != Precision::I32 &&
  34                 layer->insData[UNSQUEEZE_INDEXES].lock()->getTensorDesc().getPrecision() != Precision::FP32)
  35                 THROW_IE_EXCEPTION << layer->name << " Incorrect 'indices_to_squeeze' input precision. Only FP32 and I32 are supported!";
  36
  37             addConfig(layer, { { ConfLayout::PLN, false, 0 }, { ConfLayout::ANY, true } }, { { ConfLayout::PLN, false, 0 } });
  38         } catch (InferenceEngine::details::InferenceEngineException &ex) {
  39             errorMsg = ex.what();
  40         }
  41     }
  42
  43     StatusCode execute(std::vector<Blob::Ptr>& inputs, std::vector<Blob::Ptr>& outputs, ResponseDesc *resp) noexcept override {
  44         InferenceEngine::SizeVector data_dims = inputs[UNSQUEEZE_DATA]->getTensorDesc().getDims();
  45         InferenceEngine::SizeVector idx_dims = inputs[UNSQUEEZE_INDEXES]->getTensorDesc().getDims();
  46
  47         switch (inputs[UNSQUEEZE_INDEXES]->precision()) {
  48         case Precision::FP32: {
  49             float *idx_data = inputs[UNSQUEEZE_INDEXES]->cbuffer().as<float *>() +
  50                               inputs[UNSQUEEZE_INDEXES]->getTensorDesc().getBlockingDesc().getOffsetPadding();
  51
  52             size_t max = data_dims.size();
  53             for (size_t i = 0; i < idx_dims[0]; i++) {
  54                 size_t axis = static_cast<size_t>(idx_data[i]);
  55                 if (axis > max) max = axis;
  56             }
  57             max++;
  58
  59             if ((idx_dims[0] + data_dims.size()) < max) {
  60                 if (resp) {
  61                     std::string errorMsg = "Indices_to_set for unsqueeze layer is out of tensor dimension";
  62                     errorMsg.copy(resp->msg, sizeof(resp->msg) - 1);
  63                 }
  64                 return PARAMETER_MISMATCH;
  65             }
  66         }
  67         break;
  68         case Precision::I32: {
  69             int32_t *idx_data = inputs[UNSQUEEZE_INDEXES]->cbuffer().as<int32_t *>() +
  70                                 inputs[UNSQUEEZE_INDEXES]->getTensorDesc().getBlockingDesc().getOffsetPadding();
  71             size_t max = data_dims.size();
  72             for (size_t i = 0; i < idx_dims[0]; i++) {
  73                 size_t axis = static_cast<size_t>(idx_data[i]);
  74                 if (axis > max) max = axis;
  75             }
  76             max++;
  77
  78             if ((idx_dims[0] + data_dims.size()) < max) {
  79                 if (resp) {
  80                     std::string errorMsg = "Indices_to_set for unsqueeze layer is out of tensor dimension";
  81                     errorMsg.copy(resp->msg, sizeof(resp->msg) - 1);
  82                 }
  83                 return PARAMETER_MISMATCH;
  84             }
  85         }
  86         break;
  87         default:
  88             if (resp) {
  89                 std::string errorMsg = "Incorrect 'indices_to_set' input precision. Only FP32 and I32 are supported!";
  90                 errorMsg.copy(resp->msg, sizeof(resp->msg) - 1);
  91             }
  92             return GENERAL_ERROR;
  93         }
  94
  95         return OK;
  96     }
  97
  98 private:
  99     const size_t UNSQUEEZE_DATA = 0;
 100     const size_t UNSQUEEZE_INDEXES = 1;
 101
 102     SizeVector data_dims;
 103     SizeVector idx_dims;
 104 };
 105
 106 REG_FACTORY_FOR(ImplFactory<UnsqueezeImpl>, Unsqueeze);
 107
 108 }  // namespace Cpu
 109 }  // namespace Extensions
 110 }  // namespace InferenceEngine