wehub-resource-sync
59 lines
1.9 KiB
C++
59 lines
1.9 KiB
C++
//
|
|
// ReIndexOnnxIfAlias.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/05.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <MNN/MNNDefine.h>
|
|
#include "../PostTreatUtils.hpp"
|
|
#include "../Global.hpp"
|
|
#include "../SubGraphComplete.hpp"
|
|
using namespace MNN;
|
|
class ReIndexOnnxIfAlias : public PostConverter {
|
|
public:
|
|
virtual bool onExecute(std::unique_ptr<MNN::NetT>& net) const override {
|
|
if (net->sourceType != NetSource_ONNX) {
|
|
return true;
|
|
}
|
|
auto* ctx = Global<MNN::Express::OptimizeContext>::Get();
|
|
auto updateNameFromNewSubGraph = [=](const std::string& graphName, std::string& aliasName) {
|
|
for (auto gOld : ctx->subgraphs) {
|
|
if (gOld->name != graphName) {
|
|
continue;
|
|
}
|
|
int idx = -1;
|
|
for (int i = 0; i < gOld->outputs.size(); ++i) {
|
|
if (gOld->tensors[gOld->outputs[i]] == aliasName) {
|
|
idx = i;
|
|
break;
|
|
}
|
|
}
|
|
if (idx < 0) {
|
|
break;
|
|
}
|
|
for (auto gNew : ctx->completed_subgraphs) {
|
|
if (gNew->name != graphName) {
|
|
continue;
|
|
}
|
|
aliasName = gNew->tensors[gNew->outputs[idx]];
|
|
break;
|
|
}
|
|
}
|
|
};
|
|
for (auto& op : net->oplists) {
|
|
if (op->type != OpType_If) {
|
|
continue;
|
|
}
|
|
auto param = op->main.AsIfParam();
|
|
for (auto& pair : param->aliases_outputs) {
|
|
updateNameFromNewSubGraph(param->then_graph, pair->data[0]);
|
|
updateNameFromNewSubGraph(param->else_graph, pair->data[1]);
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
};
|
|
static PostConverterRegister<ReIndexOnnxIfAlias> __l("ReIndexOnnxIfAlias");
|