109 lines
3.2 KiB
C++
109 lines
3.2 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#include "graph/task.h"
|
|
|
|
#include "archive/value_archive.h"
|
|
#include "core/graph.h"
|
|
#include "core/operator.h"
|
|
#include "graph/common.h"
|
|
|
|
namespace mmdeploy::graph {
|
|
|
|
static int GetDepth(const Value& input) {
|
|
if (input.is_array() && input.size() > 0) {
|
|
return GetDepth(input[0]) + 1;
|
|
}
|
|
return input.is_array();
|
|
}
|
|
|
|
// all args are array of the same length
|
|
static size_t GetBatchSize(const Value& args) {
|
|
size_t batch_size = 0;
|
|
for (const auto& x : args) {
|
|
if (x.is_array()) {
|
|
if (!batch_size) {
|
|
batch_size = x.size();
|
|
} else if (batch_size != x.size()) {
|
|
return 0;
|
|
}
|
|
} else {
|
|
return 0;
|
|
}
|
|
}
|
|
return batch_size;
|
|
}
|
|
|
|
Task::Task(const Value& cfg) : BaseNode(cfg) {
|
|
auto module = CreateFromRegistry<Module>(cfg, "module");
|
|
if (!module) {
|
|
MMDEPLOY_ERROR("failed to create task: {}", cfg);
|
|
throw_exception(eFail);
|
|
}
|
|
module_ = std::move(module).value();
|
|
name_ = cfg.value("name", string{});
|
|
is_batched_ = cfg.value("is_batched", false);
|
|
is_thread_safe_ = cfg.value("is_thread_safe", false);
|
|
}
|
|
|
|
void Task::Build(TaskGraph& graph) {
|
|
auto handle = graph.Add([this](Context& ctx) -> Result<void> {
|
|
auto args = ctx.pop().array();
|
|
auto rets = Value::Array{};
|
|
auto batch_size = GetBatchSize(args);
|
|
// MMDEPLOY_ERROR("name: {}, is_batched: {}, INPUT batch_size: {}", name_, is_batched_,
|
|
// batch_size);
|
|
if (!is_batched_ && batch_size) {
|
|
rets.resize(outputs_.size(), Value::kArray);
|
|
if (!is_thread_safe_) {
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
Value sample = Value::kArray;
|
|
for (const auto& a : args) {
|
|
sample.push_back(a[i]);
|
|
}
|
|
OUTCOME_TRY(auto ret, module_->Process(sample));
|
|
for (int j = 0; j < ret.size(); ++j) {
|
|
rets[j].push_back(std::move(ret[j]));
|
|
}
|
|
}
|
|
} else {
|
|
std::vector<std::function<Result<Value>()>> tasks;
|
|
tasks.reserve(batch_size);
|
|
OUTCOME_TRY(auto batch_args, DistribAA(args));
|
|
for (int sample_id = 0; sample_id < batch_size; ++sample_id) {
|
|
tasks.emplace_back([&, sample_id]() -> Result<Value> {
|
|
return module_->Process(batch_args[sample_id]);
|
|
});
|
|
}
|
|
auto batch_rets = ctx.Execute(tasks);
|
|
for (auto& batch_ret : batch_rets) {
|
|
OUTCOME_TRY(auto ret, std::move(batch_ret));
|
|
for (int j = 0; j < rets.size(); ++j) {
|
|
rets[j].push_back(std::move(ret[j]));
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
OUTCOME_TRY(auto&& tmp, module_->Process(args));
|
|
rets = std::move(tmp).array();
|
|
}
|
|
ctx.push(std::move(rets));
|
|
// MMDEPLOY_ERROR("name: {}, is_batched: {}, OUTPUT batch_size: {}", name_, is_batched_,
|
|
// GetBatchSize(rets));
|
|
return success();
|
|
});
|
|
handle->set_name(name_);
|
|
}
|
|
|
|
class TaskNodeCreator : public Creator<Node> {
|
|
public:
|
|
const char* GetName() const override { return "Task"; }
|
|
int GetVersion() const override { return 0; }
|
|
std::unique_ptr<Node> Create(const Value& value) override {
|
|
return std::make_unique<Task>(value);
|
|
}
|
|
};
|
|
|
|
REGISTER_MODULE(Node, TaskNodeCreator);
|
|
|
|
} // namespace mmdeploy::graph
|