wehub-resource-sync
54 lines
2.0 KiB
C++
54 lines
2.0 KiB
C++
//
|
|
// MergeReluToBinaryOp.hpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/05.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "../PostTreatUtils.hpp"
|
|
using namespace MNN;
|
|
|
|
class MergeReluToBinaryOp : public PostConverter {
|
|
public:
|
|
virtual bool onExecute(std::unique_ptr<MNN::NetT>& net) const override {
|
|
// Merge Layer
|
|
std::vector<MNN::OpT*> readyToDelete;
|
|
for (auto iter = net->oplists.begin(); iter != net->oplists.end(); iter++) {
|
|
MNN::OpT& currentOp = *(iter->get());
|
|
if (currentOp.type != MNN::OpType_BinaryOp || currentOp.inputIndexes.size() != 2) {
|
|
continue;
|
|
}
|
|
DCHECK(currentOp.outputIndexes.size() == 1) << "Binary output ERROR!";
|
|
|
|
// merge Relu/Relu6 to Binary
|
|
std::vector<MNN::OpT*> nextOp = PostTreatUtils::_findOpByInputIndex(currentOp.outputIndexes[0], net.get());
|
|
while (1) {
|
|
if (nextOp.size() != 1) {
|
|
break;
|
|
}
|
|
const int nextOutputIndex = nextOp[0]->outputIndexes[0];
|
|
|
|
bool nextRelu = (nextOp[0]->type == MNN::OpType_ReLU && nextOp[0]->main.AsRelu()->slope == 0.0f);
|
|
|
|
if (PostTreatUtils::_isSingleInputOutput(nextOp[0]) && nextRelu) {
|
|
//LOG(INFO) << "Merge " << nextOp[0]->name.c_str()<< " into Binary: ";
|
|
// currentOp.name.c_str();
|
|
currentOp.main.AsBinaryOp()->activationType = 1;
|
|
currentOp.outputIndexes[0] = nextOp[0]->outputIndexes[0];
|
|
readyToDelete.push_back(nextOp[0]);
|
|
nextOp = PostTreatUtils::_findOpByInputIndex(nextOutputIndex, net.get());
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
for (auto op : readyToDelete) {
|
|
PostTreatUtils::_removeOpInNet(op, net.get());
|
|
}
|
|
return true;
|
|
}
|
|
};
|
|
|
|
static PostConverterRegister<MergeReluToBinaryOp> __l("MergeReluToBinaryOp");
|