wehub-resource-sync
121 lines
3.5 KiB
C++
121 lines
3.5 KiB
C++
//
|
|
// TFGraphResolverHelpers.hpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2020/06/13.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#ifndef TF_GRAPH_RESOLVER_HELPERS_HPP_
|
|
#define TF_GRAPH_RESOLVER_HELPERS_HPP_
|
|
|
|
#include "TFGraphResolver.hpp"
|
|
|
|
#include <vector>
|
|
#include <queue>
|
|
#include <unordered_map>
|
|
#include <unordered_set>
|
|
|
|
#include "MNN_generated.h"
|
|
#include "../compression/quantization.hpp"
|
|
|
|
inline bool IsControlInput(const std::string& name) {
|
|
return name.size() && name.substr(0, 1) == "^";
|
|
}
|
|
|
|
inline bool IsInWhileLoop(const std::string& name,
|
|
const std::unordered_map<std::string,
|
|
std::string>& types) {
|
|
return types.at(name) == "WhileLoop";
|
|
}
|
|
|
|
inline std::vector<std::string> RSplitString(const std::string& name,
|
|
const std::string& sp) {
|
|
std::vector<std::string> splits;
|
|
size_t pos = name.rfind(sp);
|
|
if (pos != std::string::npos) {
|
|
splits.push_back(name.substr(0, pos));
|
|
splits.push_back(name.substr(pos + 1));
|
|
} else {
|
|
splits.push_back(name);
|
|
}
|
|
return std::move(splits);
|
|
}
|
|
|
|
inline bool IsControlFlowNode(const TFNode* node) {
|
|
static std::unordered_set<std::string> control_flow_ops{
|
|
"Enter", "Merge", "LoopCond", "Switch", "NextIteration", "Exit",
|
|
};
|
|
return control_flow_ops.count(node->op);
|
|
}
|
|
|
|
template <typename T>
|
|
inline T NodeAttr(const TFNode* node, const std::string& attr_name) {
|
|
MNN_ERROR("This function should not be called.\n");
|
|
return T();
|
|
}
|
|
|
|
template <>
|
|
inline std::string NodeAttr<std::string>(const TFNode* node,
|
|
const std::string& attr_name) {
|
|
if (!node->node_def->attr().count(attr_name)) {
|
|
MNN_ERROR("Can not find attribute named %s.\n", attr_name.c_str());
|
|
}
|
|
return node->node_def->attr().at(attr_name).s();
|
|
}
|
|
|
|
inline void EraseInput(TFNode* node, TFEdge* edge) {
|
|
auto it = std::remove_if(node->inputs.begin(), node->inputs.end(),
|
|
[edge](const TFEdge* input) {
|
|
return input == edge;
|
|
});
|
|
node->inputs.erase(it, node->inputs.end());
|
|
}
|
|
|
|
inline void AddInput(TFNode* node, TFEdge* edge) {
|
|
node->inputs.push_back(edge);
|
|
}
|
|
|
|
inline void EraseOutput(TFNode* node, TFEdge* edge) {
|
|
auto it = std::remove_if(node->outputs.begin(), node->outputs.end(),
|
|
[edge](const TFEdge* output) {
|
|
return output == edge;
|
|
});
|
|
node->outputs.erase(it, node->outputs.end());
|
|
}
|
|
|
|
inline void AddOutput(TFNode* node, TFEdge* edge) {
|
|
node->outputs.push_back(edge);
|
|
}
|
|
|
|
template <typename NodeT, typename StopFunc>
|
|
inline std::vector<NodeT*> ReverseVisit(const std::vector<NodeT*>& final_nodes,
|
|
StopFunc stop_fn) {
|
|
std::vector<NodeT*> nodes;
|
|
std::unordered_set<NodeT*> visited;
|
|
std::queue<NodeT*> queue;
|
|
for (NodeT* node : final_nodes) {
|
|
queue.push(node);
|
|
visited.insert(node);
|
|
}
|
|
|
|
while (!queue.empty()) {
|
|
NodeT* node = queue.front();
|
|
queue.pop();
|
|
|
|
if (stop_fn(node)) {
|
|
continue;
|
|
}
|
|
nodes.push_back(node);
|
|
for (auto* edge : node->inputs) {
|
|
NodeT* start = edge->start;
|
|
if (visited.insert(start).second) {
|
|
queue.push(start);
|
|
}
|
|
}
|
|
}
|
|
return std::move(nodes);
|
|
}
|
|
|
|
#endif // TF_GRAPH_RESOLVER_HELPERS_HPP_
|