wehub-resource-sync
81 lines
2.4 KiB
C++
81 lines
2.4 KiB
C++
//
|
|
// tfOpConverter.hpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/01/31.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#ifndef TFtfOpCONVERTER_HPP
|
|
#define TFtfOpCONVERTER_HPP
|
|
#include <map>
|
|
#include <string>
|
|
#include <vector>
|
|
#include "MNN_generated.h"
|
|
#include "TmpGraph.hpp"
|
|
#include "graph.pb.h"
|
|
|
|
#include "logkit.h"
|
|
|
|
// The base class for tensorflow op converter
|
|
class tfOpConverter {
|
|
friend class tfOpConverterSuit;
|
|
|
|
public:
|
|
virtual void run(MNN::OpT *dstOp, TmpNode *srcNode) = 0;
|
|
virtual MNN::OpParameter type() = 0;
|
|
virtual MNN::OpType opType() = 0;
|
|
tfOpConverter() {
|
|
}
|
|
virtual ~tfOpConverter() {
|
|
}
|
|
|
|
static void convertTensorToBlob(MNN::BlobT* dst, const ::tensorflow::TensorProto& tensor);
|
|
|
|
private:
|
|
};
|
|
|
|
// converter factory
|
|
class tfOpConverterSuit {
|
|
public:
|
|
static tfOpConverterSuit *get();
|
|
void insert(tfOpConverter *t, const char *name);
|
|
|
|
tfOpConverter *search(const std::string &name);
|
|
tfOpConverterSuit() {
|
|
}
|
|
~tfOpConverterSuit();
|
|
|
|
private:
|
|
static tfOpConverterSuit *global;
|
|
std::map<std::string, tfOpConverter *> mTests;
|
|
};
|
|
|
|
template <class T>
|
|
class tfOpConverterRegister {
|
|
public:
|
|
tfOpConverterRegister(const char *claim) {
|
|
T *test = new T;
|
|
tfOpConverterSuit *ts = tfOpConverterSuit::get();
|
|
ts->insert(test, claim);
|
|
}
|
|
~tfOpConverterRegister() {
|
|
}
|
|
};
|
|
|
|
#define DECLARE_OP_CONVERTER(name) \
|
|
class name : public tfOpConverter { \
|
|
public: \
|
|
virtual void run(MNN::OpT *dstOp, TmpNode *srcNode); \
|
|
name() { \
|
|
} \
|
|
virtual ~name() { \
|
|
} \
|
|
virtual MNN::OpType opType(); \
|
|
virtual MNN::OpParameter type(); \
|
|
}
|
|
|
|
#define REGISTER_CONVERTER(name, opType) static tfOpConverterRegister<name> _Convert_##opType(#opType)
|
|
|
|
#endif // TFtfOpCONVERTER_HPP
|