wehub-resource-sync
108 lines
4.0 KiB
C++
108 lines
4.0 KiB
C++
//
|
|
// OnnxPad.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2020/07/09.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <map>
|
|
#include <string>
|
|
#include <vector>
|
|
#include "MNN_generated.h"
|
|
#include "OnnxExtraManager.hpp"
|
|
#include "logkit.h"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
class OnnxPadTransform : public OnnxExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
auto op = expr->get();
|
|
auto opName = op->name()->str();
|
|
|
|
PadValueMode mode = CONSTANT;
|
|
VARP padsVar;
|
|
bool padsFromInput = true;
|
|
|
|
auto info = op->main_as_Extra();
|
|
if (nullptr != info->attr()) {
|
|
for (int i = 0; i < info->attr()->size(); ++i) {
|
|
const auto attr = info->attr()->GetAs<Attribute>(i);
|
|
const auto attributeName = attr->key()->str();
|
|
if (attributeName == "mode") {
|
|
const std::map<std::string, PadValueMode> padValueModeMap = {
|
|
{"constant", CONSTANT}, {"reflect", REFLECT}, {"edge", EDGE}
|
|
};
|
|
auto modeStr = attr->s()->str();
|
|
if (padValueModeMap.find(modeStr) == padValueModeMap.end()) {
|
|
LOG(ERROR) << "MNN only support ['constant', 'reflect'] Pad mode";
|
|
return nullptr;
|
|
}
|
|
mode = padValueModeMap.at(modeStr);
|
|
} else if (attributeName == "pads") {
|
|
padsFromInput = false;
|
|
auto padList = attr->list()->i();
|
|
int size = padList->size();
|
|
std::vector<int> pads(size);
|
|
for (int s = 0; s < size / 2; ++s) {
|
|
pads[s * 2] = padList->Get(s);
|
|
pads[s * 2 + 1] = padList->Get(s + size / 2);
|
|
}
|
|
padsVar = _Const(pads.data(), {(int)pads.size()}, NCHW, halide_type_of<int>());
|
|
}
|
|
}
|
|
}
|
|
if (padsFromInput) {
|
|
if (inputs.size() == 1) {
|
|
LOG(ERROR) << "MNN need pad value in attr or other node";
|
|
return nullptr;
|
|
}
|
|
/* Pytorch's Pad exported by Onnx (opset_version=11) have complicated subgraph.
|
|
Pad values in input node is [before_pads, after_pads], which is not same as MNN pads order.
|
|
Example: pad2d, onnx: [left, upper, right, bottom], MNN: [left, right, upper, bottom]
|
|
So we need this order converting subgraph (all const, not affect inference speed).
|
|
*/
|
|
padsVar = _Reshape(_Transpose(_Reshape(inputs[1], {2, -1}), {1, 0}), {-1});
|
|
}
|
|
std::unique_ptr<OpT> pad(new OpT);
|
|
pad->type = OpType_Padding;
|
|
pad->main.type = OpParameter_PadParam;
|
|
pad->main.value = new PadParamT;
|
|
switch (mode) {
|
|
case CONSTANT:
|
|
pad->main.AsPadParam()->mode = MNN::PadValueMode_CONSTANT;
|
|
break;
|
|
case SYMMETRIC:
|
|
pad->main.AsPadParam()->mode = MNN::PadValueMode_SYMMETRIC;
|
|
break;
|
|
case REFLECT:
|
|
pad->main.AsPadParam()->mode = MNN::PadValueMode_REFLECT;
|
|
break;
|
|
case EDGE:
|
|
pad->main.AsPadParam()->mode = MNN::PadValueMode_EDGE;
|
|
break;
|
|
default:
|
|
pad->main.AsPadParam()->mode = MNN::PadValueMode_CONSTANT;
|
|
break;
|
|
}
|
|
std::vector<VARP> newInputs{inputs[0], padsVar};
|
|
if (inputs.size() > 2) {
|
|
newInputs.emplace_back(inputs[2]);
|
|
}
|
|
auto res = Expr::create(pad.get(), newInputs);
|
|
res->setName(opName);
|
|
return res;
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
OnnxExtraManager::get()->insert("Pad", std::shared_ptr<OnnxExtraManager::Transform>(new OnnxPadTransform));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|