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

   1 // Copyright (C) 2018-2019 Intel Corporation
   2 // SPDX-License-Identifier: Apache-2.0
   3 //
   4
   5 #pragma once
   6
   7 #include <description_buffer.hpp>
   8 #include "ie_built_in_impl.hpp"
   9 #include <shape_infer/ie_reshaper.hpp>
  10 #include <ie_layers.h>
  11 #include <map>
  12 #include <memory>
  13 #include <string>
  14 #include <vector>
  15
  16 namespace InferenceEngine {
  17 namespace ShapeInfer {
  18
  19 /**
  20  *@brief Implementation of Shape inference for DetectionOutput layer
  21  */
  22 class TensorIteratorShapeProp : public BuiltInShapeInferImpl {
  23 public:
  24     explicit TensorIteratorShapeProp(const std::string& type) : BuiltInShapeInferImpl(type) {}
  25
  26     void setOriginalLayer(const CNNLayer *layer) {
  27         auto ti = dynamic_cast<const TensorIterator*>(layer);
  28         if (!ti)
  29             THROW_IE_EXCEPTION << "Error during shape infer. Original layer is not TensorIterator.";
  30         _original_ti = ti;
  31     }
  32
  33     void inferShapesImpl(const std::vector<Blob::CPtr>& inBlobs,
  34                          const std::map<std::string, std::string>& params,
  35                          const std::map<std::string, Blob::Ptr>& blobs,
  36                          std::vector<SizeVector>& outShapes) override {
  37         LayerParams lp{};
  38         TensorIterator ti(lp);
  39         ti.params = params;
  40         ti.type = _type;
  41         ti.body = _original_ti->body;
  42         ti.back_edges = _original_ti->back_edges;
  43         ti.input_port_map = _original_ti->input_port_map;
  44         ti.output_port_map = _original_ti->output_port_map;
  45         validate(&ti, inBlobs, params, blobs);
  46
  47         // TODO: make util function to calculate num of iteration
  48         int num_iteration = 1;
  49
  50         // Prepare input shapes for internal body
  51         std::map<std::string, std::vector<size_t>> newInShapes;
  52         for (auto &port_map : ti.input_port_map) {
  53             int ext_port = port_map.from;
  54             int int_port = port_map.to;
  55             auto int_name = ti.body.inputs[int_port]->name;
  56
  57             auto shape = inShapes[ext_port];
  58             if (port_map.axis != -1) {
  59                 int size = shape[port_map.axis];
  60                 int start = port_map.start < 0
  61                         ? port_map.start + size + 1
  62                         : port_map.start;
  63                 int end = port_map.end < 0
  64                         ? port_map.end + size + 1
  65                         : port_map.end;
  66
  67                 num_iteration = std::abs(end - start) / std::abs(port_map.stride);
  68
  69                 // port with iterating through. Change dimension with iteration
  70                 shape[port_map.axis] = port_map.part_size;
  71             }
  72
  73             newInShapes[int_name] = shape;
  74         }
  75
  76         // Body shape infer
  77         _body_reshaper = std::make_shared<Reshaper>(_original_ti->body.inputs);
  78         _body_reshaper->runNoApply(newInShapes);
  79
  80         outShapes.resize(ti.output_port_map.size());
  81         for (auto &port_map : ti.output_port_map) {
  82             int ext_port = port_map.from;
  83             int int_port = port_map.to;
  84             auto &int_out_data = ti.body.outputs[int_port];
  85             auto shape = _body_reshaper->getResultShapeFor(int_out_data);
  86
  87             if (port_map.axis != -1) {
  88                 // port with iterating through. Change dimension with iteration
  89                 shape[port_map.axis] *= num_iteration;
  90             }
  91
  92             outShapes[ext_port] = shape;
  93         }
  94     }
  95
  96     void apply() {
  97         if (!_body_reshaper)
  98             THROW_IE_EXCEPTION << "Request of apply reshape results while shape infer was not finished";
  99         _body_reshaper->apply();
 100     }
 101
 102
 103 private:
 104     const TensorIterator* _original_ti;
 105     std::shared_ptr<Reshaper> _body_reshaper;
 106 };
 107
 108 }  // namespace ShapeInfer
 109 }  // namespace InferenceEngine