wehub-resource-sync
104 lines
2.4 KiB
C++
104 lines
2.4 KiB
C++
//
|
|
// TFGraphResolver.hpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2020/06/13.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#ifndef TF_GRAPH_RESOLVER_HPP_
|
|
#define TF_GRAPH_RESOLVER_HPP_
|
|
|
|
#include <vector>
|
|
#include <unordered_map>
|
|
|
|
#include "MNN/MNNDefine.h"
|
|
#include "graph.pb.h"
|
|
#include "MNN_generated.h"
|
|
|
|
#include "../compression/quantization.hpp"
|
|
|
|
struct TFNode;
|
|
typedef tensorflow::NodeDef NodeDef;
|
|
|
|
struct TFEdge {
|
|
std::string name;
|
|
TFNode* start;
|
|
TFNode* end;
|
|
};
|
|
|
|
struct TFNode {
|
|
const NodeDef* node_def;
|
|
std::string name;
|
|
std::string op;
|
|
std::vector<TFEdge*> inputs;
|
|
std::vector<TFEdge*> outputs;
|
|
};
|
|
|
|
class TFGraph {
|
|
public:
|
|
TFGraph() = default;
|
|
virtual ~TFGraph() = default;
|
|
|
|
void AddNode(const NodeDef* node);
|
|
void Finalize();
|
|
|
|
std::unique_ptr<MNN::SubGraphProtoT> ToProto() const;
|
|
|
|
private:
|
|
friend class TFGraphResolver;
|
|
std::string name_ = "main";
|
|
|
|
std::vector<std::unique_ptr<TFNode>> nodes_;
|
|
std::vector<std::unique_ptr<TFEdge>> edges_;
|
|
std::vector<std::unique_ptr<NodeDef>> allocated_nodes_;
|
|
// Output nodes.
|
|
std::vector<TFNode*> final_nodes_;
|
|
};
|
|
|
|
class TFGraphResolver {
|
|
public:
|
|
explicit TFGraphResolver(const tensorflow::GraphDef& graph_def);
|
|
virtual ~TFGraphResolver() = default;
|
|
|
|
TFGraph* graph(const int graph_index);
|
|
|
|
const TFGraph* graph(const int graph_index) const;
|
|
|
|
int graph_size() const { return graphs_.size(); }
|
|
|
|
private:
|
|
TFGraph* main_graph();
|
|
|
|
void ResolveQuantization(TFGraph* graph,
|
|
const compression::Quantization& int8_calibration);
|
|
|
|
std::unique_ptr<TFNode> BuildQuantOrDequantNode(
|
|
const std::string& name,
|
|
const std::string& op,
|
|
const int& nbit,
|
|
const std::vector<float>& scales,
|
|
const float& zero_point,
|
|
const float& clamp_min,
|
|
const float& clamp_max,
|
|
const MNN::Compression::LayerQuantizeParams_QuantMethod& method);
|
|
|
|
std::unique_ptr<TFEdge> BuildEdge(
|
|
const std::string& name, TFNode* start, TFNode* end);
|
|
|
|
private:
|
|
std::vector<std::unique_ptr<TFGraph>> graphs_;
|
|
|
|
struct StringComp {
|
|
bool operator()(const std::string& lhs, const std::string& rhs) const {
|
|
if (lhs.size() != rhs.size()) {
|
|
return lhs.size() > rhs.size();
|
|
} else {
|
|
return lhs > rhs;
|
|
}
|
|
}
|
|
};
|
|
};
|
|
|
|
#endif // TF_GRAPH_RESOLVER_HPP_
|