1
0
mirror of https://github.com/open-mmlab/mmdeploy.git synced 2025-01-14 08:09:43 +08:00
mmdeploy/csrc/net/net_module.cpp
lzhangzz 640aa03538
Support Windows ()
* minor changes

* support windows

* fix GCC build

* fix lint

* reformat

* fix Windows build

* fix GCC build

* search backend ops for onnxruntime

* fix lint

* fix lint

* code clean-up

* code clean-up

* fix clang build

* fix trt support

* fix cmake for ncnn

* fix cmake for openvino

* fix SDK Python API

* handle ops for other backends (ncnn, trt)

* handle SDK Python API library location

* robustify linkage

* fix cuda

* minor fix for openvino & ncnn

* use CMAKE_CUDA_ARCHITECTURES if set

* fix cuda preprocessor

* fix misc

* fix pplnn & pplcv, drop support for pplcv<0.6.0

* robustify cmake

* update build.md ()

* build dynamic modules as module library & fix demo (partially)

* fix candidate path for mmdeploy_python

* move "enable CUDA" to cmake config for demo

* refine demo cmake

* add comment

* fix ubuntu build

* revert docs/en/build.md

* fix C API

* fix lint

* Windows build doc ()

* check in docs related to mmdeploy build on windows

* update build guide on windows platform

* update build guide on windows platform

* make path of thirdparty libraries consistent

* make path consistency

* correct build command for custom ops

* correct build command for sdk

* update sdk build instructions

* update doc

* correct build command

* fix lint

* correct build command and fix lint

Co-authored-by: lvhan <lvhan@pjlab.org>

* trailing whitespace ()

* minor fix

* fix sr sdk model

* fix type deduction

* fix cudaFree after driver shutting down

* update ppl.cv installation warning ()

* fix device allocator threshold & fix lint

* update doc ()

* update ppl.cv installation warning

* missing 'git clone'

Co-authored-by: chenxin <chenxin2@sensetime.com>
Co-authored-by: zhangli <zhangli@sensetime.com>
Co-authored-by: lvhan028 <lvhan_028@163.com>
Co-authored-by: lvhan <lvhan@pjlab.org>
2022-02-24 20:08:44 +08:00

238 lines
7.1 KiB
C++

// Copyright (c) OpenMMLab. All rights reserved.
#include "net_module.h"
#include <thread>
#include "archive/value_archive.h"
#include "core/logger.h"
#include "core/model.h"
#include "core/module.h"
#include "core/net.h"
#include "core/registry.h"
#include "core/utils/formatter.h"
#include "core/utils/scope_counter.h"
#include "experimental/module_adapter.h"
using std::string;
using std::vector;
namespace mmdeploy {
struct NetModule::Impl {
using Input = std::map<std::string, Tensor>;
using Output = std::map<std::string, Tensor>;
explicit Impl(const Value& args) {
MMDEPLOY_DEBUG("Net Module cfg: {}", args);
auto init = [&]() -> Result<void> {
auto name = args["name"].get<std::string>();
auto& context = args["context"];
auto model = context["model"].get<Model>();
OUTCOME_TRY(auto config, model.GetModelConfig(name));
device_ = context.value("device", Device{"cpu"});
stream_ = context.value("stream", Stream::GetDefault(device_));
auto creator = Registry<Net>::Get().GetCreator(config.backend);
if (!creator) {
MMDEPLOY_ERROR("Net backend not found: {}", config.backend);
return Status(eEntryNotFound);
}
auto net_cfg = args;
net_cfg["context"].update({{"device", device_}, {"stream", stream_}});
net_ = creator->Create(net_cfg);
if (!net_) {
return Status(eFail);
}
OUTCOME_TRY(InitializeInputTensors(args));
OUTCOME_TRY(InitializeOutputTensors(args));
return success();
};
init().value();
}
Result<void> InitializeInputTensors(const Value& args) {
auto inputs = args.value<Value>("input_map", ValueType::kObject);
for (auto it = inputs.begin(); it != inputs.end(); ++it) {
input_mapping_.insert({(*it).get<std::string>(), it.key()});
}
OUTCOME_TRY(inputs_, net_->GetInputTensors());
for (const auto& t : inputs_) {
input_mapping_.insert({t.name(), t.name()});
}
return success();
}
Result<void> InitializeOutputTensors(const Value& args) {
auto outputs = args.value<Value>("output_map", ValueType::kObject);
for (auto it = outputs.begin(); it != outputs.end(); ++it) {
output_mapping_.insert({(*it).get<std::string>(), it.key()});
}
OUTCOME_TRY(outputs_, net_->GetOutputTensors());
for (const auto& t : outputs_) {
output_mapping_.insert({t.name(), t.name()});
}
return success();
}
Result<TensorShape> InferInputShape(const vector<Tensor>& input) {
auto batch_size = input.size();
auto& exemplar = input.front();
auto shape = exemplar.shape();
if (batch_size == 1) {
return shape;
}
if (shape[0] != 1) {
MMDEPLOY_ERROR("unsupported shape for batch assemble: {}", shape);
return Status(eNotSupported);
}
for (int i = 1; i < input.size(); ++i) {
auto& sample = input[i];
if (sample.shape() != shape) {
MMDEPLOY_ERROR("shapes are not consistent across the batch");
return Status(eNotSupported);
}
}
shape[0] = static_cast<int64_t>(batch_size);
return shape;
}
Result<vector<TensorShape> > InferInputShape(const vector<vector<Tensor> >& inputs) {
vector<TensorShape> shapes;
shapes.reserve(inputs.size());
for (const auto& input : inputs) {
OUTCOME_TRY(auto shape, InferInputShape(input));
shapes.push_back(std::move(shape));
}
return shapes;
}
Result<std::vector<Output> > Forward(const std::vector<Input>& input) {
// auto t0 = std::chrono::high_resolution_clock::now();
//
auto batch_size = static_cast<int>(input.size());
std::vector<std::vector<Tensor> > input_samples;
input_samples.reserve(inputs_.size());
for (const auto& t : inputs_) {
auto name = input_mapping_.at(t.name());
std::vector<Tensor> tmp;
tmp.reserve(input.size());
for (int i = 0; i < input.size(); ++i) {
auto& sample = input[i];
if (auto it = sample.find(name); it != sample.end()) {
tmp.push_back(it->second);
} else {
MMDEPLOY_ERROR("sample {} missing key {}", i, name);
return Status(eInvalidArgument);
}
}
input_samples.push_back(std::move(tmp));
}
// 1. calculate input shape
OUTCOME_TRY(auto input_shapes, InferInputShape(input_samples));
// 2. call backend's reshape
OUTCOME_TRY(net_->Reshape(input_shapes));
// 3. fill input tensor
for (int i = 0; i < inputs_.size(); ++i) {
auto& src = input_samples[i];
auto& dst = inputs_[i];
if (dst.shape() != input_shapes[i]) {
MMDEPLOY_ERROR("inconsistent input shape, expect {}, got {}", input_shapes[i], dst.shape());
return Status(eFail);
}
if (src.size() > 1) {
for (int j = 0; j < src.size(); ++j) {
auto slice = dst.Slice(j);
OUTCOME_TRY(src[j].CopyTo(slice, stream_));
}
} else {
OUTCOME_TRY(src[0].CopyTo(dst, stream_));
}
}
// 5. forward
OUTCOME_TRY(net_->Forward());
vector<Output> output(batch_size);
for (const auto& t : outputs_) {
auto name = output_mapping_.at(t.name());
auto desc = t.desc();
desc.device = device_;
Tensor tmp(desc);
if (tmp.size()) {
OUTCOME_TRY(t.CopyTo(tmp, stream_));
} else {
MMDEPLOY_WARN("copy skipped due to zero sized tensor");
}
if (output.size() > 1) {
for (int i = 0; i < output.size(); ++i) {
output[i].emplace(name, tmp.Slice(i));
}
} else {
output[0].emplace(name, std::move(tmp));
}
}
return output;
}
Device device_;
Stream stream_;
std::unique_ptr<Net> net_;
Span<Tensor> inputs_;
Span<Tensor> outputs_;
// outer scope to model input names
std::map<std::string, std::string> input_mapping_;
// outer scope to model output names
std::map<std::string, std::string> output_mapping_;
};
NetModule::~NetModule() = default;
NetModule::NetModule(const Value& args) : impl_(std::make_unique<Impl>(args)) {}
Result<Value> NetModule::operator()(const Value& input) {
auto filter = [](const Value& sample) {
Impl::Input tensors;
for (auto it = sample.begin(); it != sample.end(); ++it) {
if (it->is_any<Tensor>()) {
tensors.insert({it.key(), it->get<Tensor>()});
}
}
return tensors;
};
std::vector<Impl::Input> batch;
if (input.is_array()) {
batch.reserve(input.size());
for (const auto& sample : input) {
batch.push_back(filter(sample));
}
} else if (input.is_object()) {
batch.push_back(filter(input));
} else {
return Status(eNotSupported);
}
OUTCOME_TRY(auto batch_output, impl_->Forward(batch));
if (input.is_array()) {
return to_value(batch_output);
} else {
return to_value(batch_output.at(0));
}
}
class NetModuleCreator : public Creator<Module> {
public:
const char* GetName() const override { return "Net"; }
int GetVersion() const override { return 0; }
std::unique_ptr<Module> Create(const Value& value) override {
return CreateTask(NetModule{value});
}
};
REGISTER_MODULE(Module, NetModuleCreator);
} // namespace mmdeploy