wehub-resource-sync
86 lines
3.1 KiB
C++
86 lines
3.1 KiB
C++
//
|
|
// torchConverter.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2020/11/13.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "MNN_generated.h"
|
|
#include "torchOpConverter.hpp"
|
|
#include "torchOptimize.hpp"
|
|
#include <torch/csrc/jit/passes/freeze_module.h>
|
|
#if !defined(_MSC_VER)
|
|
#include <dlfcn.h>
|
|
#endif
|
|
|
|
void loadCustomOp(std::string customTorchOps) {
|
|
if (customTorchOps.empty()) {
|
|
return;
|
|
}
|
|
#if !defined(_MSC_VER)
|
|
constexpr char delimiter = ';';
|
|
std::string::size_type lastPos = customTorchOps.find_first_not_of(delimiter, 0);
|
|
std::string::size_type pos = customTorchOps.find_first_of(delimiter, lastPos);
|
|
while (std::string::npos != pos || std::string::npos != lastPos) {
|
|
auto custom_lib = customTorchOps.substr(lastPos, pos - lastPos);
|
|
dlopen(custom_lib.c_str(), RTLD_NOW | RTLD_LOCAL);
|
|
lastPos = customTorchOps.find_first_not_of(delimiter, pos);
|
|
pos = customTorchOps.find_first_of(delimiter, lastPos);
|
|
}
|
|
#endif
|
|
}
|
|
|
|
MNN_PUBLIC int torch2MNNNet(const std::string inputModel, const std::string bizCode,
|
|
std::unique_ptr<MNN::NetT>& netT, std::string customTorchOps) {
|
|
loadCustomOp(customTorchOps);
|
|
// Deserialize the ScriptModule from a file, set to eval mode and freeze
|
|
c10::Device device("cpu");
|
|
torch::jit::Module module;
|
|
try {
|
|
module = torch::jit::load(inputModel.c_str(), device);
|
|
} catch (std::exception e) {
|
|
MNN_ERROR("[ERROR] TorchScript model can't load. Please using `torch.jit.script` or `torch.jit.trace` save model.\n");
|
|
return 1;
|
|
}
|
|
auto graph = torch::jit::torchOptPass(module);
|
|
std::unique_ptr<TorchScope> scope(new TorchScope(netT.get()));
|
|
for (const auto input : graph->inputs()) {
|
|
auto type = input->type()->cast<at::TensorType>();
|
|
if (!type) {
|
|
continue;
|
|
}
|
|
auto scalarType = type->scalarType().value_or(at::ScalarType::Float);
|
|
auto inputName = input->debugName();
|
|
scope->declareTensor(inputName);
|
|
MNN::OpT* MNNOp = new MNN::OpT;
|
|
MNNOp->name = inputName;
|
|
MNNOp->type = MNN::OpType_Input;
|
|
MNNOp->main.type = MNN::OpParameter_Input;
|
|
auto param = new MNN::InputT;
|
|
param->dtype = ScalarType2Dtype(scalarType);
|
|
param->dformat = MNN::MNN_DATA_FORMAT_NCHW;
|
|
MNNOp->main.value = param;
|
|
netT->oplists.emplace_back(MNNOp);
|
|
MNNOp->outputIndexes.push_back(scope->lookupTensor(inputName));
|
|
}
|
|
for (const auto &output : graph->outputs()) {
|
|
netT->outputName.push_back(output->debugName());
|
|
}
|
|
for (const auto &node : graph->nodes()) {
|
|
const auto& kind = node->kind();
|
|
bool isOutputNode = false;
|
|
for (const auto output : node->outputs()) {
|
|
isOutputNode |= std::find(netT->outputName.begin(), netT->outputName.end(), output->debugName()) != netT->outputName.end();
|
|
}
|
|
// python prim ops
|
|
if (!isOutputNode && kind.is_prim() && scope->dealPrime(node)) {
|
|
continue;
|
|
}
|
|
scope->buildMNNOp(node);
|
|
}
|
|
netT->sourceType = MNN::NetSource_TORCH;
|
|
netT->bizCode = bizCode;
|
|
return 0;
|
|
}
|