wehub-resource-sync
111 lines
4.4 KiB
C++
111 lines
4.4 KiB
C++
//
|
|
// RemoveCopy.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/05.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "../PostTreatUtils.hpp"
|
|
#include "config.hpp"
|
|
#include "../Global.hpp"
|
|
class RemoveCopy : public PostConverter {
|
|
public:
|
|
virtual bool onExecute(std::unique_ptr<MNN::NetT>& net) const override {
|
|
std::set<std::string> netOutputNames;
|
|
for (auto& t : net->outputName) {
|
|
netOutputNames.insert(t);
|
|
}
|
|
for (auto iter = net->oplists.begin(); iter != net->oplists.end(); iter++) {
|
|
auto& op = *iter;
|
|
if (op->type == MNN::OpType_Input) {
|
|
for (auto o : op->outputIndexes) {
|
|
netOutputNames.insert(net->tensorName[o]);
|
|
}
|
|
}
|
|
}
|
|
std::map<int, std::unique_ptr<MNN::TensorDescribeT>> desmap;
|
|
for (auto&& iter : net->extraTensorDescribe) {
|
|
desmap.insert(std::make_pair(iter->index, std::move(iter)));
|
|
}
|
|
auto config = Global<modelConfig>::Get();
|
|
for (auto iter = net->oplists.begin(); iter != net->oplists.end();) {
|
|
auto& op = *iter;
|
|
if (op->type != MNN::OpType_Identity || op->inputIndexes.size() != op->outputIndexes.size()) {
|
|
iter++;
|
|
continue;
|
|
}
|
|
|
|
bool hasOutputName = false;
|
|
for (auto o : op->outputIndexes) {
|
|
if (netOutputNames.find(net->tensorName[o]) != netOutputNames.end()) {
|
|
hasOutputName = true;
|
|
break;
|
|
}
|
|
}
|
|
bool hasOutputFromInput = false;
|
|
for (auto o : op->inputIndexes) {
|
|
if (netOutputNames.find(net->tensorName[o]) != netOutputNames.end()) {
|
|
hasOutputFromInput = true;
|
|
break;
|
|
}
|
|
}
|
|
if (hasOutputFromInput && hasOutputName) {
|
|
iter++;
|
|
continue;
|
|
}
|
|
auto originInput = op->inputIndexes;
|
|
auto originOutputs = op->outputIndexes;
|
|
MNN_ASSERT(originInput.size() == originOutputs.size());
|
|
if (hasOutputName) {
|
|
bool valid = true;
|
|
for (int i=0; i<op->inputIndexes.size(); ++i) {
|
|
auto o = op->outputIndexes[i];
|
|
auto originInput = op->inputIndexes[i];
|
|
if (netOutputNames.find(net->tensorName[o]) != netOutputNames.end()) {
|
|
if (netOutputNames.find(net->tensorName[originInput]) != netOutputNames.end()) {
|
|
valid = false;
|
|
break;
|
|
}
|
|
auto originName = net->tensorName[originInput];
|
|
net->tensorName[originInput] = net->tensorName[o];
|
|
net->tensorName[o] = originName;
|
|
}
|
|
}
|
|
if (!valid) {
|
|
continue;
|
|
}
|
|
}
|
|
|
|
std::map<int, int> replaceIndexes;
|
|
for (int i=0; i<op->inputIndexes.size();++i) {
|
|
replaceIndexes.insert(std::make_pair(op->outputIndexes[i], op->inputIndexes[i]));
|
|
}
|
|
for (auto& replaceIter : replaceIndexes) {
|
|
auto desIter = desmap.find(replaceIter.first);
|
|
if (desIter != desmap.end()) {
|
|
desIter->second->index = replaceIter.second;
|
|
desIter->second->name = net->tensorName[replaceIter.second];
|
|
desmap[replaceIter.second] = std::move(desIter->second);
|
|
desmap.erase(desIter);
|
|
}
|
|
}
|
|
for (auto subIter = net->oplists.begin(); subIter != net->oplists.end(); subIter++) {
|
|
auto& subOp = *subIter;
|
|
for (int v = 0; v < subOp->inputIndexes.size(); ++v) {
|
|
if (replaceIndexes.find(subOp->inputIndexes[v]) != replaceIndexes.end()) {
|
|
subOp->inputIndexes[v] = replaceIndexes[subOp->inputIndexes[v]];
|
|
}
|
|
}
|
|
}
|
|
iter = net->oplists.erase(iter);
|
|
}
|
|
net->extraTensorDescribe.clear();
|
|
for (auto&& iter : desmap) {
|
|
net->extraTensorDescribe.emplace_back(std::move(iter.second));
|
|
}
|
|
return true;
|
|
}
|
|
};
|
|
static PostConverterRegister<RemoveCopy> __l("RemoveCopy");
|