File size: 2,680 Bytes
be903e2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | // Tencent is pleased to support the open source community by making ncnn available.
//
// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.
#include "constant_unpooling.h"
#include <unordered_map>
#include <unordered_set>
namespace pnnx {
void ConstantUnpooling(std::shared_ptr<torch::jit::Graph>& graph, torch::jit::Block* block, std::unordered_set<torch::jit::Node*>& constants)
{
for (auto it = block->nodes().begin(); it != block->nodes().end();)
{
auto node = *it;
// node may be moved to a different block so advance iterator now
++it;
if (!node->blocks().empty())
{
// Traverse sub-blocks.
for (auto block : node->blocks())
{
ConstantUnpooling(graph, block, constants);
}
continue;
}
for (int i = 0; i < (int)node->inputs().size(); i++)
{
const auto& in = node->input(i);
if (in->node()->kind() != c10::prim::Constant)
continue;
// input constant node
if (constants.find(in->node()) == constants.end())
{
constants.insert(in->node());
continue;
}
torch::jit::WithInsertPoint guard(node);
std::unordered_map<torch::jit::Value*, torch::jit::Value*> value_map;
auto value_map_func = [&](torch::jit::Value* v) {
return value_map.at(v);
};
// graph->setInsertPoint(node);
auto* new_constant_node = graph->insertNode(graph->createClone(in->node(), value_map_func, false));
// fprintf(stderr, "new_constant_node %s\n", new_constant_node->outputs()[0]->debugName().c_str());
// create new constant node
node->replaceInput(i, new_constant_node->outputs()[0]);
}
}
}
void constant_unpooling(std::shared_ptr<torch::jit::Graph>& graph)
{
std::unordered_set<torch::jit::Node*> constants;
ConstantUnpooling(graph, graph->block(), constants);
}
} // namespace pnnx
|