wehub-resource-sync
124 lines
3.9 KiB
C++
124 lines
3.9 KiB
C++
//
|
|
// ReduceTorch.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2021/05/13.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <stdio.h>
|
|
#include "torchOpConverter.hpp"
|
|
|
|
DECLARE_OP_CONVERTER(ReduceTorch);
|
|
|
|
MNN::OpType ReduceTorch::opType() {
|
|
return MNN::OpType_Reduction;
|
|
}
|
|
|
|
MNN::OpParameter ReduceTorch::type() {
|
|
return MNN::OpParameter_ReductionParam;
|
|
}
|
|
|
|
std::vector<int> ReduceTorch::inputTensorIdx() {
|
|
return {0};
|
|
}
|
|
|
|
void ReduceTorch::run(MNN::OpT* dstOp, const torch::jit::Node* node, TorchScope* scope) {
|
|
static std::map<std::string, MNN::ReductionType> gMaps{
|
|
{"sum_reduce", MNN::ReductionType_SUM},
|
|
{"mean", MNN::ReductionType_MEAN},
|
|
{"max_reduce", MNN::ReductionType_MAXIMUM},
|
|
{"amax", MNN::ReductionType_MAXIMUM},
|
|
{"min_reduce", MNN::ReductionType_MINIMUM},
|
|
{"amin", MNN::ReductionType_MINIMUM},
|
|
{"prod", MNN::ReductionType_PROD},
|
|
{"all", MNN::ReductionType_ALL},
|
|
{"any", MNN::ReductionType_ANY},
|
|
};
|
|
const auto inputs = node->inputs();
|
|
auto param = new MNN::ReductionParamT;
|
|
std::string opType = getRealOpType(node);
|
|
param->operation = gMaps[opType];
|
|
|
|
if (opType == "sum_reduce" || opType == "mean" || opType == "amax" || opType == "amin") {
|
|
const auto dims = getValue<std::vector<int64_t>>(inputs[1]);
|
|
for (int i : dims) {
|
|
param->dim.push_back(i);
|
|
}
|
|
param->keepDims = getValue<bool>(inputs[2]);
|
|
} else {
|
|
if (inputs[1]->type()->kind() == c10::TypeKind::IntType) {
|
|
const auto dim = getValue<int64_t>(inputs[1]);
|
|
param->dim.push_back(dim);
|
|
} else {
|
|
const auto dims = getValue<std::vector<int64_t>>(inputs[1]);
|
|
for (auto dim : dims) {
|
|
param->dim.push_back(dim);
|
|
}
|
|
}
|
|
}
|
|
if (dstOp->outputIndexes.size() > 1) {
|
|
int realOutput = dstOp->outputIndexes[0];
|
|
dstOp->outputIndexes.clear();
|
|
dstOp->outputIndexes.push_back(realOutput);
|
|
}
|
|
dstOp->main.value = param;
|
|
}
|
|
|
|
REGISTER_CONVERTER(ReduceTorch, sum_reduce);
|
|
REGISTER_CONVERTER(ReduceTorch, mean);
|
|
REGISTER_CONVERTER(ReduceTorch, max_reduce);
|
|
REGISTER_CONVERTER(ReduceTorch, min_reduce);
|
|
REGISTER_CONVERTER(ReduceTorch, prod);
|
|
REGISTER_CONVERTER(ReduceTorch, all);
|
|
REGISTER_CONVERTER(ReduceTorch, any);
|
|
REGISTER_CONVERTER(ReduceTorch, amin);
|
|
REGISTER_CONVERTER(ReduceTorch, amax);
|
|
|
|
DECLARE_OP_CONVERTER(NormTorch);
|
|
|
|
MNN::OpType NormTorch::opType() {
|
|
return MNN::OpType_Extra;
|
|
}
|
|
MNN::OpParameter NormTorch::type() {
|
|
return MNN::OpParameter_Extra;
|
|
}
|
|
std::vector<int> NormTorch::inputTensorIdx() {
|
|
return {0};
|
|
}
|
|
|
|
void NormTorch::run(MNN::OpT* dstOp, const torch::jit::Node* node, TorchScope* scope) {
|
|
auto extra = new MNN::ExtraT;
|
|
dstOp->main.value = extra;
|
|
extra->engine = "Torch";
|
|
extra->type = "norm";
|
|
auto type = getRealOpType(node);
|
|
extra->attr.resize(3);
|
|
extra->attr[0].reset(new MNN::AttributeT);
|
|
extra->attr[0]->key = "ord";
|
|
extra->attr[1].reset(new MNN::AttributeT);
|
|
extra->attr[1]->key = "dim";
|
|
extra->attr[2].reset(new MNN::AttributeT);
|
|
extra->attr[2]->key = "keepDim";
|
|
if (type == "frobenius_norm") {
|
|
extra->attr[0]->i = 2;
|
|
auto dims = getValue<std::vector<int64_t>>(node->input(1));
|
|
extra->attr[1]->i = dims[0];
|
|
extra->attr[2]->i = getValue<bool>(node->input(2));
|
|
} else {
|
|
auto ord = node->input(1);
|
|
auto kind = ord->type()->kind();
|
|
if (kind == c10::TypeKind::FloatType) {
|
|
extra->attr[0]->i = getValue<double>(node->input(1));
|
|
} else {
|
|
extra->attr[0]->i = getValue<int64_t>(node->input(1));
|
|
}
|
|
auto dims = getValue<std::vector<int64_t>>(node->input(2));
|
|
extra->attr[1]->i = dims[0];
|
|
extra->attr[2]->i = getValue<bool>(node->input(3));
|
|
}
|
|
}
|
|
|
|
REGISTER_CONVERTER(NormTorch, norm);
|
|
REGISTER_CONVERTER(NormTorch, frobenius_norm);
|