wehub-resource-sync
79 lines
2.6 KiB
C++
79 lines
2.6 KiB
C++
//
|
|
// CastTorch.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2021/07/22.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <stdio.h>
|
|
#include "torchOpConverter.hpp"
|
|
|
|
DECLARE_OP_CONVERTER(CastTorch);
|
|
|
|
MNN::OpType CastTorch::opType() {
|
|
return MNN::OpType_Cast;
|
|
}
|
|
MNN::OpParameter CastTorch::type() {
|
|
return MNN::OpParameter_CastParam;
|
|
}
|
|
std::vector<int> CastTorch::inputTensorIdx() {
|
|
return {0};
|
|
}
|
|
|
|
void CastTorch::run(MNN::OpT* dstOp, const torch::jit::Node* node, TorchScope* scope) {
|
|
auto param = new MNN::CastParamT;
|
|
std::string opType = getRealOpType(node);
|
|
if (opType == "to" || opType == "as_tensor") {
|
|
/*
|
|
scalar_type_to_pytorch_type = [
|
|
torch.uint8, # 0
|
|
torch.int8, # 1
|
|
torch.short, # 2
|
|
torch.int, # 3
|
|
torch.int64, # 4
|
|
torch.half, # 5
|
|
torch.float, # 6
|
|
torch.double, # 7
|
|
torch.complex32, # 8
|
|
torch.complex64, # 9
|
|
torch.complex128, # 10
|
|
torch.bool, # 11
|
|
]
|
|
*/
|
|
static std::vector<MNN::DataType> gMaps {
|
|
MNN::DataType_DT_UINT8, MNN::DataType_DT_INT8, MNN::DataType_DT_INT16,
|
|
MNN::DataType_DT_INT32, MNN::DataType_DT_INT64, MNN::DataType_DT_HALF,
|
|
MNN::DataType_DT_FLOAT, MNN::DataType_DT_DOUBLE,
|
|
MNN::DataType_DT_INVALID, MNN::DataType_DT_INVALID, MNN::DataType_DT_INVALID,
|
|
MNN::DataType_DT_BOOL,
|
|
};
|
|
param->dstT = gMaps[getValue<int64_t>(node->input(1))];
|
|
} else if (opType == "type_as") {
|
|
auto type = node->input(1)->type()->cast<at::TensorType>();
|
|
if (type) {
|
|
auto scalarType = type->scalarType().value_or(at::ScalarType::Float);
|
|
param->dstT = ScalarType2Dtype(scalarType);
|
|
} else {
|
|
param->dstT = MNN::DataType_DT_FLOAT;
|
|
}
|
|
} else {
|
|
static std::map<std::string, MNN::DataType> gMaps{
|
|
{"Int", MNN::DataType_DT_INT32}, {"IntImplicit", MNN::DataType_DT_INT32},
|
|
{"Bool", MNN::DataType_DT_BOOL},
|
|
{"Float", MNN::DataType_DT_FLOAT}, {"FloatImplicit", MNN::DataType_DT_FLOAT},
|
|
};
|
|
param->dstT = gMaps[opType];
|
|
}
|
|
dstOp->main.value = param;
|
|
}
|
|
|
|
REGISTER_CONVERTER(CastTorch, Int);
|
|
REGISTER_CONVERTER(CastTorch, IntImplicit);
|
|
REGISTER_CONVERTER(CastTorch, Bool);
|
|
REGISTER_CONVERTER(CastTorch, Float);
|
|
REGISTER_CONVERTER(CastTorch, FloatImplicit);
|
|
REGISTER_CONVERTER(CastTorch, to);
|
|
REGISTER_CONVERTER(CastTorch, type_as);
|
|
REGISTER_CONVERTER(CastTorch, as_tensor);
|