wehub-resource-sync
114 lines
4.2 KiB
C++
114 lines
4.2 KiB
C++
//
|
|
// ReIndexTensor.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/05.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <MNN/MNNDefine.h>
|
|
#include <set>
|
|
#include <sstream>
|
|
#include "../PostTreatUtils.hpp"
|
|
using namespace MNN;
|
|
class ReIndexTensor : public PostConverter {
|
|
public:
|
|
virtual bool onExecute(std::unique_ptr<MNN::NetT>& net) const override {
|
|
auto& mNet = net;
|
|
std::map<std::string, int> tensorNameIdx;
|
|
std::map<int, int> usefulTensorIndexMap;
|
|
std::vector<std::string> usefulTensorName;
|
|
// extraTensorDescribe reindex
|
|
for (int i = 0; i < mNet->tensorName.size(); i++) {
|
|
tensorNameIdx.insert(std::make_pair(mNet->tensorName[i], i));
|
|
}
|
|
for (int i = 0; i < mNet->extraTensorDescribe.size(); i++) {
|
|
auto name = mNet->extraTensorDescribe[i]->name;
|
|
auto iter = tensorNameIdx.find(name);
|
|
if (iter == tensorNameIdx.end()) {
|
|
mNet->extraTensorDescribe[i]->index = -1;
|
|
} else {
|
|
mNet->extraTensorDescribe[i]->index = iter->second;
|
|
}
|
|
}
|
|
|
|
std::vector<bool> tensorValid(mNet->tensorName.size(), false);
|
|
for (auto& op : mNet->oplists) {
|
|
for (auto index : op->inputIndexes) {
|
|
if (index < 0) {
|
|
continue; // optional input, ignore it
|
|
}
|
|
tensorValid[index] = true;
|
|
}
|
|
for (auto index : op->outputIndexes) {
|
|
tensorValid[index] = true;
|
|
}
|
|
}
|
|
|
|
for (int i = 0; i < tensorValid.size(); ++i) {
|
|
if (tensorValid[i]) {
|
|
usefulTensorIndexMap.insert(std::make_pair(i, usefulTensorName.size()));
|
|
usefulTensorName.push_back(mNet->tensorName[i]);
|
|
}
|
|
}
|
|
|
|
// Re index
|
|
for (auto& op : mNet->oplists) {
|
|
for (int i = 0; i < op->inputIndexes.size(); ++i) {
|
|
if (op->inputIndexes[i] < 0) {
|
|
continue;
|
|
}
|
|
auto iter = usefulTensorIndexMap.find(op->inputIndexes[i]);
|
|
DCHECK(iter != usefulTensorIndexMap.end()) << "ERROR";
|
|
op->inputIndexes[i] = iter->second;
|
|
}
|
|
for (int i = 0; i < op->outputIndexes.size(); ++i) {
|
|
auto iter = usefulTensorIndexMap.find(op->outputIndexes[i]);
|
|
DCHECK(iter != usefulTensorIndexMap.end()) << "ERROR";
|
|
op->outputIndexes[i] = iter->second;
|
|
}
|
|
}
|
|
|
|
mNet->tensorName = usefulTensorName;
|
|
for (auto iter = mNet->extraTensorDescribe.begin(); iter != mNet->extraTensorDescribe.end();) {
|
|
auto index = (*iter)->index;
|
|
if (usefulTensorIndexMap.find(index) == usefulTensorIndexMap.end()) {
|
|
iter = mNet->extraTensorDescribe.erase(iter);
|
|
continue;
|
|
}
|
|
(*iter)->index = usefulTensorIndexMap.find(index)->second;
|
|
iter++;
|
|
}
|
|
// Check dup name and modify
|
|
std::set<std::string> names;
|
|
std::set<std::string> tensorNames;
|
|
for (int i = 0; i < mNet->oplists.size(); ++i) {
|
|
auto& op = mNet->oplists[i];
|
|
auto opName = op->name;
|
|
if (opName.empty() || names.find(opName) != names.end()) {
|
|
std::ostringstream defaultName;
|
|
defaultName << EnumNameOpType(op->type);
|
|
defaultName << i;
|
|
op->name = defaultName.str();
|
|
#ifdef DEBUG
|
|
MNN_PRINT("%d op name is empty or dup, set to %s\n", i, op->name.c_str());
|
|
#endif
|
|
opName = op->name;
|
|
}
|
|
names.insert(opName);
|
|
for (auto output : op->outputIndexes) {
|
|
auto origin = net->tensorName[output];
|
|
if (origin.empty() || tensorNames.find(origin) != tensorNames.end()) {
|
|
std::ostringstream defaultName;
|
|
defaultName << output;
|
|
origin = defaultName.str();
|
|
net->tensorName[output] = origin;
|
|
}
|
|
tensorNames.insert(origin);
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
};
|
|
static PostConverterRegister<ReIndexTensor> __l("ReIndexTensor");
|