wehub-resource-sync
42 lines
1.1 KiB
C++
42 lines
1.1 KiB
C++
//
|
|
// SplitTorch.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2021/08/05.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <stdio.h>
|
|
#include "torchOpConverter.hpp"
|
|
|
|
DECLARE_OP_CONVERTER(SplitTorch);
|
|
|
|
MNN::OpType SplitTorch::opType() {
|
|
return MNN::OpType_Slice;
|
|
}
|
|
MNN::OpParameter SplitTorch::type() {
|
|
return MNN::OpParameter_Slice;
|
|
}
|
|
std::vector<int> SplitTorch::inputTensorIdx() {
|
|
return {0};
|
|
}
|
|
|
|
void SplitTorch::run(MNN::OpT* dstOp, const torch::jit::Node* node, TorchScope* scope) {
|
|
auto param = new MNN::SliceT;
|
|
auto kind = node->input(1)->type()->kind();
|
|
if (kind == c10::TypeKind::IntType) {
|
|
param->slicePoints.push_back(getValue<int64_t>(node->input(1)));
|
|
} else {
|
|
auto split_sizes = getValue<std::vector<int64_t>>(node->input(1));
|
|
for (auto i : split_sizes) {
|
|
param->slicePoints.push_back(i);
|
|
}
|
|
}
|
|
param->axis = getValue<int64_t>(node->input(2));
|
|
// TORCH Split: param is INT: split_size; param is [INT]: [split_size_0, split_size_1, ...]
|
|
param->sourceType = MNN::NetSource_TORCH;
|
|
dstOp->main.value = param;
|
|
}
|
|
|
|
REGISTER_CONVERTER(SplitTorch, split);
|