wehub-resource-sync
127 lines
4.8 KiB
C++
127 lines
4.8 KiB
C++
//
|
|
// ConstantTorch.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2021/05/12.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <stdio.h>
|
|
#include "torchOpConverter.hpp"
|
|
|
|
DECLARE_OP_CONVERTER(ConstantTorch);
|
|
|
|
MNN::OpType ConstantTorch::opType() {
|
|
return MNN::OpType_Const;
|
|
}
|
|
MNN::OpParameter ConstantTorch::type() {
|
|
return MNN::OpParameter_Blob;
|
|
}
|
|
std::vector<int> ConstantTorch::inputTensorIdx() {
|
|
return {};
|
|
}
|
|
|
|
void ConstantTorch::run(MNN::OpT* dstOp, const torch::jit::Node* node, TorchScope* scope) {
|
|
auto param = new MNN::BlobT;
|
|
const auto output = node->output();
|
|
const std::string& type = output->type()->str();
|
|
if (type == "None") {
|
|
return;
|
|
}
|
|
if (getRealOpType(node) == "Uninitialized" || node->attributeNames().empty()) {
|
|
param->dataType = MNN::DataType_DT_FLOAT;
|
|
param->float32s = {};
|
|
dstOp->main.value = param;
|
|
return;
|
|
}
|
|
auto attr = node->attributeNames()[0];
|
|
auto kind = node->kindOf(attr);
|
|
switch (kind) {
|
|
case torch::jit::AttributeKind::f:
|
|
param->dataType = MNN::DataType_DT_FLOAT;
|
|
param->float32s.push_back(node->f(attr));
|
|
break;
|
|
case torch::jit::AttributeKind::i:
|
|
param->dataType = MNN::DataType_DT_INT32;
|
|
// node->i is int64_t
|
|
param->int32s.push_back(std::min(node->i(attr), static_cast<int64_t>(std::numeric_limits<int>::max())));
|
|
break;
|
|
case torch::jit::AttributeKind::s:
|
|
param->dataType = MNN::DataType_DT_STRING;
|
|
param->strings.push_back(node->s(attr));
|
|
break;
|
|
case torch::jit::AttributeKind::ival: {
|
|
param->dataType = MNN::DataType_DT_INT32;
|
|
const auto int64s = getValue<std::vector<int64_t>>(output);
|
|
param->int32s.resize(int64s.size());
|
|
param->dims.push_back(int64s.size());
|
|
for (int i = 0; i < int64s.size(); i++) {
|
|
param->int32s[i] = int64s[i];
|
|
}
|
|
break;
|
|
}
|
|
case torch::jit::AttributeKind::t: {
|
|
const auto tensor = getValue<at::Tensor>(output);
|
|
auto scalarType = tensor.scalar_type();
|
|
switch (scalarType) {
|
|
case at::ScalarType::Byte:
|
|
param->dataType = MNN::DataType_DT_UINT8;
|
|
param->uint8s = std::move(getValue<uint8_t>(output, param->dims));
|
|
break;
|
|
case at::ScalarType::Char:
|
|
param->dataType = MNN::DataType_DT_INT8;
|
|
param->int8s = std::move(getValue<int8_t>(output, param->dims));
|
|
break;
|
|
case at::ScalarType::Int:
|
|
param->dataType = MNN::DataType_DT_INT32;
|
|
param->int32s = std::move(getValue<int32_t>(output, param->dims));
|
|
break;
|
|
case at::ScalarType::Long: {
|
|
param->dataType = MNN::DataType_DT_INT32;
|
|
const auto int64s = std::move(getValue<int64_t>(output, param->dims));
|
|
param->int32s.resize(int64s.size());
|
|
for (int i = 0; i < int64s.size(); i++) {
|
|
param->int32s[i] = int64s[i];
|
|
}
|
|
break;
|
|
}
|
|
case at::ScalarType::Float:
|
|
param->dataType = MNN::DataType_DT_FLOAT;
|
|
param->float32s = std::move(getValue<float>(output, param->dims));
|
|
break;
|
|
case at::ScalarType::Double: {
|
|
param->dataType = MNN::DataType_DT_FLOAT;
|
|
const auto doubles = getValue<double>(output, param->dims);
|
|
param->float32s.resize(doubles.size());
|
|
for (int i = 0; i < doubles.size(); i++) {
|
|
param->float32s[i] = doubles[i];
|
|
}
|
|
break;
|
|
}
|
|
case at::ScalarType::Bool:
|
|
param->dataType = MNN::DataType_DT_INT32;
|
|
param->int32s = std::move(getValue<int32_t>(output, param->dims));
|
|
if (param->dims.empty() && param->int32s.empty()) {
|
|
param->int32s.push_back(0);
|
|
param->dims.push_back(1);
|
|
}
|
|
break;
|
|
case at::ScalarType::BFloat16:
|
|
case at::ScalarType::Short:
|
|
case at::ScalarType::Half:
|
|
default:
|
|
MNN_ASSERT(false);
|
|
break;
|
|
}
|
|
break;
|
|
}
|
|
default:
|
|
MNN_ASSERT(false);
|
|
return;
|
|
}
|
|
dstOp->main.value = param;
|
|
}
|
|
|
|
REGISTER_CONVERTER(ConstantTorch, Constant);
|
|
REGISTER_CONVERTER(ConstantTorch, Uninitialized);
|