wehub-resource-sync
107 lines
6.1 KiB
C++
107 lines
6.1 KiB
C++
//
|
|
// TransformShuffleChannel.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/06.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "../PostTreatUtils.hpp"
|
|
using namespace MNN;
|
|
class TransformShuffleChannel : public PostConverter {
|
|
public:
|
|
virtual bool onExecute(std::unique_ptr<MNN::NetT>& net) const override {
|
|
for (auto iter = net->oplists.begin(); iter != net->oplists.end();) {
|
|
auto op = iter->get();
|
|
if (op->type == OpType_Plugin) {
|
|
auto plugin = op->main.AsPlugin();
|
|
if (plugin->type == "ShuffleChannel") {
|
|
int currentTensorCount = (int)net->tensorName.size();
|
|
std::unique_ptr<OpT> convertTo(new OpT);
|
|
convertTo->type = OpType_ConvertTensor;
|
|
convertTo->main.type = OpParameter_TensorConvertInfo;
|
|
convertTo->main.value = new TensorConvertInfoT;
|
|
convertTo->main.AsTensorConvertInfo()->source = MNN_DATA_FORMAT_NC4HW4;
|
|
convertTo->main.AsTensorConvertInfo()->dest = MNN_DATA_FORMAT_NHWC;
|
|
convertTo->inputIndexes = op->inputIndexes;
|
|
convertTo->name = op->name + "_ConvertToNHWC";
|
|
convertTo->outputIndexes = {currentTensorCount + 0};
|
|
net->tensorName.emplace_back(convertTo->name);
|
|
|
|
auto group = plugin->attr[0]->tensor->int32s[0];
|
|
std::unique_ptr<OpT> reshape(new OpT);
|
|
reshape->type = OpType_Reshape;
|
|
reshape->name = op->name + "_Reshape";
|
|
reshape->main.value = new ReshapeT;
|
|
reshape->main.type = OpParameter_Reshape;
|
|
reshape->main.AsReshape()->dimType = MNN_DATA_FORMAT_NHWC;
|
|
reshape->main.AsReshape()->dims = {0, 0, 0, group, -1};
|
|
reshape->inputIndexes = {currentTensorCount + 0};
|
|
reshape->outputIndexes = {currentTensorCount + 1};
|
|
net->tensorName.emplace_back(reshape->name);
|
|
|
|
std::unique_ptr<OpT> constOp(new OpT);
|
|
constOp->type = OpType_Const;
|
|
auto blob = new BlobT;
|
|
blob->int32s = {0, 1, 2, 4, 3};
|
|
blob->dataFormat = MNN_DATA_FORMAT_NHWC;
|
|
blob->dataType = DataType_DT_INT32;
|
|
blob->dims = {5};
|
|
constOp->main.value = blob;
|
|
constOp->main.type = OpParameter_Blob;
|
|
constOp->name = op->name + "_Const";
|
|
constOp->outputIndexes = {currentTensorCount + 2};
|
|
net->tensorName.emplace_back(constOp->name);
|
|
|
|
std::unique_ptr<OpT> permute(new OpT);
|
|
permute->type = OpType_Transpose;
|
|
permute->name = op->name + "_Transpose";
|
|
permute->main.value = new TransposeT;
|
|
permute->main.type = OpParameter_Transpose;
|
|
permute->main.AsTranspose()->Tperm = DataType_DT_INT32;
|
|
permute->inputIndexes = {currentTensorCount + 1, currentTensorCount + 2};
|
|
permute->outputIndexes = {currentTensorCount + 3};
|
|
net->tensorName.emplace_back(permute->name);
|
|
|
|
std::unique_ptr<OpT> reshapeR(new OpT);
|
|
reshapeR->type = OpType_Reshape;
|
|
reshapeR->name = op->name + "_ReshapeR";
|
|
reshapeR->main.value = new ReshapeT;
|
|
reshapeR->main.type = OpParameter_Reshape;
|
|
reshapeR->main.AsReshape()->dimType = MNN_DATA_FORMAT_NHWC;
|
|
reshapeR->main.AsReshape()->dims = {0, 0, 0, -1};
|
|
reshapeR->inputIndexes = {currentTensorCount + 3};
|
|
reshapeR->outputIndexes = {currentTensorCount + 4};
|
|
net->tensorName.emplace_back(reshapeR->name);
|
|
|
|
std::unique_ptr<OpT> convertFrom(new OpT);
|
|
convertFrom->type = OpType_ConvertTensor;
|
|
convertFrom->main.type = OpParameter_TensorConvertInfo;
|
|
convertFrom->main.value = new TensorConvertInfoT;
|
|
convertFrom->main.AsTensorConvertInfo()->source = MNN_DATA_FORMAT_NHWC;
|
|
convertFrom->main.AsTensorConvertInfo()->dest = MNN_DATA_FORMAT_NC4HW4;
|
|
convertFrom->inputIndexes = {currentTensorCount + 4};
|
|
convertFrom->outputIndexes = op->outputIndexes;
|
|
convertFrom->name = op->name + "_ConvertToNC4HW4";
|
|
|
|
iter = net->oplists.erase(iter);
|
|
iter = net->oplists.insert(iter, std::move(convertFrom));
|
|
iter = net->oplists.insert(iter, std::move(reshapeR));
|
|
iter = net->oplists.insert(iter, std::move(permute));
|
|
iter = net->oplists.insert(iter, std::move(constOp));
|
|
iter = net->oplists.insert(iter, std::move(reshape));
|
|
iter = net->oplists.insert(iter, std::move(convertTo));
|
|
iter++;
|
|
iter++;
|
|
iter++;
|
|
iter++;
|
|
continue;
|
|
}
|
|
}
|
|
iter++;
|
|
}
|
|
return true;
|
|
}
|
|
};
|
|
static PostConverterRegister<TransformShuffleChannel> __l("TransformShuffleChannel");
|