1 // Copyright (C) 2018-2019 Intel Corporation
2 // SPDX-License-Identifier: Apache-2.0
5 #include <vpu/frontend/frontend.hpp>
8 #include <unordered_set>
16 class RegionYoloStage final : public StageNode {
18 StagePtr cloneImpl() const override {
19 return std::make_shared<RegionYoloStage>(*this);
22 void propagateDataOrderImpl() const override {
23 IE_ASSERT(_inputEdges.size() == 1);
24 IE_ASSERT(_outputEdges.size() == 1);
26 auto output = _outputEdges[0]->output();
28 if (!attrs().get<bool>("doSoftMax")) {
29 _orderInfo.setOutput(_outputEdges[0], output->desc().dimsOrder().createMovedDim(Dim::C, 2)); // CHW
33 void getDataStridesRequirementsImpl() const override {
34 IE_ASSERT(_inputEdges.size() == 1);
35 IE_ASSERT(_outputEdges.size() == 1);
37 if (attrs().get<bool>("doSoftMax")) {
38 // Major dimension must be compact.
39 _stridesInfo.setInput(_inputEdges[0], StridesRequirement().add(2, DimStride::Compact));
43 void finalizeDataLayoutImpl() override {
46 void getBatchSupportInfoImpl() const override {
47 IE_ASSERT(_inputEdges.size() == 1);
48 IE_ASSERT(_outputEdges.size() == 1);
50 _batchInfo.setInput(_inputEdges[0], BatchSupport::Split);
51 _batchInfo.setOutput(_outputEdges[0], BatchSupport::Split);
54 void finalCheckImpl() const override {
57 void serializeParamsImpl(BlobSerializer& serializer) const override {
58 auto classes = attrs().get<int>("classes");
59 auto coords = attrs().get<int>("coords");
60 auto num = attrs().get<int>("num");
61 auto maskSize = attrs().get<int>("maskSize");
62 auto doSoftMax = attrs().get<bool>("doSoftMax");
64 serializer.append(static_cast<int32_t>(classes));
65 serializer.append(static_cast<int32_t>(coords));
66 serializer.append(static_cast<int32_t>(num));
67 serializer.append(static_cast<int32_t>(maskSize));
68 serializer.append(static_cast<int32_t>(doSoftMax));
71 void serializeDataImpl(BlobSerializer& serializer) const override {
72 IE_ASSERT(_inputEdges.size() == 1);
73 IE_ASSERT(_outputEdges.size() == 1);
74 IE_ASSERT(_tempBufferEdges.empty());
76 auto input = _inputEdges[0]->input();
77 auto output = _outputEdges[0]->output();
79 input->serializeOldBuffer(handle_from_this(), serializer);
80 output->serializeOldBuffer(handle_from_this(), serializer);
86 void FrontEnd::parseRegionYolo(
87 const Model::Ptr& model,
88 const ie::CNNLayerPtr& layer,
89 const DataVector& inputs,
90 const DataVector& outputs) {
91 IE_ASSERT(inputs.size() == 1);
92 IE_ASSERT(outputs.size() == 1);
94 auto mask = layer->GetParamAsInts("mask", {});
96 auto stage = model->addNewStage<RegionYoloStage>(
98 StageType::RegionYolo,
103 stage->attrs().set<int>("classes", layer->GetParamAsInt("classes", 20));
104 stage->attrs().set<int>("coords", layer->GetParamAsInt("coords", 4));
105 stage->attrs().set<int>("num", layer->GetParamAsInt("num", 5));
106 stage->attrs().set<int>("maskSize", static_cast<int>(mask.size()));
107 stage->attrs().set<bool>("doSoftMax", layer->GetParamAsInt("do_softmax", 1));