inference-engine/src/inference_engine/shape_infer/built-in/ie_unsqueeze_shape_infer.hpp

   1 // Copyright (C) 2019 Intel Corporation
   2 // SPDX-License-Identifier: Apache-2.0
   3 //
   4
   5 #pragma once
   6
   7 #include "ie_built_in_impl.hpp"
   8 #include <map>
   9 #include <memory>
  10 #include <string>
  11 #include <vector>
  12 #include <iostream>
  13
  14 namespace InferenceEngine {
  15 namespace ShapeInfer {
  16
  17 /**
  18  *@brief Implementation of Shape inference for Unsqueeze layer
  19  */
  20 class UnsqueezeShapeProp : public BuiltInShapeInferImpl {
  21 public:
  22     explicit UnsqueezeShapeProp(const std::string& type) : BuiltInShapeInferImpl(type) {}
  23
  24     void inferShapesImpl(const std::vector<Blob::CPtr>& inBlobs,
  25                          const std::map<std::string, std::string>& params,
  26                          const std::map<std::string, Blob::Ptr>& blobs,
  27                          std::vector<SizeVector>& outShapes) override {
  28         LayerParams lp{};
  29         UnsqueezeLayer unsqueezeLayer(lp);
  30         unsqueezeLayer.params = params;
  31         unsqueezeLayer.type = _type;
  32         validate(&unsqueezeLayer, inBlobs, params, blobs);
  33
  34         const size_t UNSQUEEZE_DATA = 0;
  35         const size_t UNSQUEEZE_INDEXES = 1;
  36
  37         SizeVector idx_dims = inBlobs[UNSQUEEZE_INDEXES]->getTensorDesc().getDims();
  38         SizeVector data_dims = inBlobs[UNSQUEEZE_DATA]->getTensorDesc().getDims();
  39         SizeVector outShape;
  40         if (idx_dims.size() > 1)
  41             THROW_IE_EXCEPTION << " Index vector should be 1 dimension";
  42         if (inBlobs[UNSQUEEZE_INDEXES]->getTensorDesc().getPrecision() != Precision::I32 &&
  43             inBlobs[UNSQUEEZE_INDEXES]->getTensorDesc().getPrecision() != Precision::FP32)
  44             THROW_IE_EXCEPTION << " Incorrect 'indices_to_squeeze' input precision. Only FP32 and I32 are supported!";
  45
  46         size_t max = data_dims.size();
  47         switch (inBlobs[UNSQUEEZE_INDEXES]->precision()) {
  48             case Precision::FP32: {
  49                 float* idx_data = inBlobs[UNSQUEEZE_INDEXES]->cbuffer().as<float*>() +
  50                                   inBlobs[UNSQUEEZE_INDEXES]->getTensorDesc().getBlockingDesc().getOffsetPadding();
  51
  52                 for (size_t i = 0; i < idx_dims[0]; i++) {
  53                     auto axis = static_cast<size_t>(idx_data[i]);
  54                     if (axis > max) max = axis;
  55                 }
  56                 max++;
  57                 if ((idx_dims[0] + data_dims.size()) < max) {
  58                     THROW_IE_EXCEPTION << "Indices_to_set for unsqueeze layer is out of tensor dimension";
  59                 }
  60                 max = inBlobs[UNSQUEEZE_INDEXES]->size() + data_dims.size();
  61                 for (size_t i = 0, j = 0, k = 0; i < max; i++) {
  62                     if (k < inBlobs[UNSQUEEZE_INDEXES]->size() && i == idx_data[k]) {
  63                         outShape.push_back(1);
  64                         k++;
  65                     } else {
  66                         outShape.push_back(data_dims[j++]);
  67                     }
  68                 }
  69             }
  70                 break;
  71             case Precision::I32: {
  72                 int32_t* idx_data = inBlobs[UNSQUEEZE_INDEXES]->cbuffer().as<int32_t*>() +
  73                                     inBlobs[UNSQUEEZE_INDEXES]->getTensorDesc().getBlockingDesc().getOffsetPadding();
  74                 max = data_dims.size();
  75                 for (size_t i = 0; i < idx_dims[0]; i++) {
  76                     auto axis = static_cast<size_t>(idx_data[i]);
  77                     if (axis > max) max = axis;
  78                 }
  79                 max++;
  80                 if ((idx_dims[0] + data_dims.size()) < max) {
  81                     THROW_IE_EXCEPTION << "Indices_to_set for unsqueeze layer is out of tensor dimension";
  82                 }
  83                 max = inBlobs[UNSQUEEZE_INDEXES]->size() + data_dims.size();
  84                 for (size_t i = 0, j = 0, k = 0; i < max; i++) {
  85                     if (k < inBlobs[UNSQUEEZE_INDEXES]->size() && i == idx_data[k]) {
  86                         outShape.push_back(1);
  87                         k++;
  88                     } else {
  89                         outShape.push_back(data_dims[j++]);
  90                     }
  91                 }
  92             }
  93             default:
  94                 THROW_IE_EXCEPTION << "Incorrect 'indices_to_set' input precision. Only FP32 and I32 are supported!";
  95         }
  96         outShapes.push_back(outShape);
  97     }
  98 };
  99
 100 }  // namespace ShapeInfer
 101 }  // namespace InferenceEngine
 102