inference-engine/src/vpu/graph_transformer/src/passes/merge_relu_and_bias.cpp

   1 // Copyright (C) 2018-2019 Intel Corporation
   2 // SPDX-License-Identifier: Apache-2.0
   3 //
   4
   5 #include <vpu/pass_manager.hpp>
   6
   7 #include <set>
   8 #include <memory>
   9
  10 #include <vpu/sw/utility.hpp>
  11
  12 namespace vpu {
  13
  14 namespace {
  15
  16 class PassImpl final : public Pass {
  17 public:
  18     explicit PassImpl(const StageBuilder::Ptr& stageBuilder) : _stageBuilder(stageBuilder) {}
  19
  20     void run(const Model::Ptr& model) override;
  21
  22 private:
  23     StageBuilder::Ptr _stageBuilder;
  24 };
  25
  26 void PassImpl::run(const Model::Ptr& model) {
  27     VPU_PROFILE(mergeReLUAndBias);
  28
  29     for (const auto& biasStage : model->getStages()) {
  30         if (biasStage == nullptr) {
  31             continue;
  32         }
  33
  34         if (biasStage->type() != StageType::Bias) {
  35             continue;
  36         }
  37
  38         if (auto reluStage = getNextStage(biasStage, {StageType::Relu, StageType::LeakyRelu})) {
  39             auto biasInput = biasStage->input(0);
  40             auto biases = biasStage->input(1);
  41
  42             auto reluOutput = reluStage->output(0);
  43
  44             auto reluStageName = reluStage->name();
  45             auto reluOrigLayer = reluStage->origLayer();
  46             auto negativeSlope = reluStage->attrs().get<float>("negativeSlope");
  47
  48             model->removeStage(biasStage);
  49             model->removeStage(reluStage);
  50
  51             _stageBuilder->addReLUStage(
  52                 model,
  53                 reluStageName,
  54                 reluOrigLayer,
  55                 negativeSlope,
  56                 biasInput,
  57                 reluOutput,
  58                 biases);
  59         }
  60     }
  61 }
  62
  63 }  // namespace
  64
  65 Pass::Ptr PassManager::mergeReLUAndBias() {
  66     return std::make_shared<PassImpl>(_stageBuilder);
  67 }
  68
  69 }  // namespace vpu