wehub-resource-sync
118 lines
3.2 KiB
C++
118 lines
3.2 KiB
C++
//
|
|
// MatMulTorch.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2021/05/10.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <stdio.h>
|
|
#include "torchOpConverter.hpp"
|
|
|
|
DECLARE_OP_CONVERTER(MatMulTorch);
|
|
|
|
MNN::OpType MatMulTorch::opType() {
|
|
return MNN::OpType_MatMul;
|
|
}
|
|
MNN::OpParameter MatMulTorch::type() {
|
|
return MNN::OpParameter_MatMul;
|
|
}
|
|
std::vector<int> MatMulTorch::inputTensorIdx() {
|
|
return {0, 1};
|
|
}
|
|
|
|
void MatMulTorch::run(MNN::OpT* dstOp, const torch::jit::Node* node, TorchScope* scope) {
|
|
auto param = new MNN::MatMulT;
|
|
std::string opType = getRealOpType(node);
|
|
if (opType == "linear") {
|
|
std::vector<int> shape;
|
|
param->bias = getValue<float>(node->input(2), shape);
|
|
param->transposeB = true;
|
|
}
|
|
dstOp->main.value = param;
|
|
}
|
|
|
|
REGISTER_CONVERTER(MatMulTorch, matmul);
|
|
REGISTER_CONVERTER(MatMulTorch, linear);
|
|
|
|
DECLARE_OP_CONVERTER(BatchMatMulTorch);
|
|
|
|
MNN::OpType BatchMatMulTorch::opType() {
|
|
return MNN::OpType_BatchMatMul;
|
|
}
|
|
MNN::OpParameter BatchMatMulTorch::type() {
|
|
return MNN::OpParameter_BatchMatMulParam;
|
|
}
|
|
std::vector<int> BatchMatMulTorch::inputTensorIdx() {
|
|
return {0, 1};
|
|
}
|
|
|
|
void BatchMatMulTorch::run(MNN::OpT* dstOp, const torch::jit::Node* node, TorchScope* scope) {
|
|
auto param = new MNN::BatchMatMulParamT;
|
|
dstOp->main.value = param;
|
|
}
|
|
|
|
REGISTER_CONVERTER(BatchMatMulTorch, bmm);
|
|
|
|
DECLARE_OP_CONVERTER(AddmmTorch);
|
|
|
|
MNN::OpType AddmmTorch::opType() {
|
|
return MNN::OpType_Extra;
|
|
}
|
|
MNN::OpParameter AddmmTorch::type() {
|
|
return MNN::OpParameter_Extra;
|
|
}
|
|
std::vector<int> AddmmTorch::inputTensorIdx() {
|
|
return {0, 1, 2};
|
|
}
|
|
|
|
void AddmmTorch::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 = getRealOpType(node);
|
|
const auto inputs = node->inputs();
|
|
const auto beta = inputs[3];
|
|
const auto alpha = inputs[4];
|
|
extra->attr.resize(2);
|
|
extra->attr[0].reset(new MNN::AttributeT);
|
|
extra->attr[0]->key = "beta";
|
|
extra->attr[0]->i = getValue<int64_t>(beta);
|
|
extra->attr[1].reset(new MNN::AttributeT);
|
|
extra->attr[1]->key = "alpha";
|
|
extra->attr[1]->i = getValue<int64_t>(alpha);
|
|
}
|
|
|
|
REGISTER_CONVERTER(AddmmTorch, addmm);
|
|
|
|
DECLARE_OP_CONVERTER(EinsumTorch);
|
|
|
|
MNN::OpType EinsumTorch::opType() {
|
|
return MNN::OpType_Extra;
|
|
}
|
|
MNN::OpParameter EinsumTorch::type() {
|
|
return MNN::OpParameter_Extra;
|
|
}
|
|
std::vector<int> EinsumTorch::inputTensorIdx() {
|
|
return {1};
|
|
}
|
|
|
|
void EinsumTorch::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 = getRealOpType(node);
|
|
const auto inputs = node->inputs();
|
|
const auto beta = inputs[3];
|
|
const auto alpha = inputs[4];
|
|
extra->attr.resize(2);
|
|
extra->attr[0].reset(new MNN::AttributeT);
|
|
extra->attr[0]->key = "beta";
|
|
extra->attr[0]->i = getValue<int64_t>(beta);
|
|
extra->attr[1].reset(new MNN::AttributeT);
|
|
extra->attr[1]->key = "alpha";
|
|
extra->attr[1]->i = getValue<int64_t>(alpha);
|
|
}
|
|
|
|
REGISTER_CONVERTER(EinsumTorch, einsum);
|