37 lines
1019 B
C++
37 lines
1019 B
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#include "graph/unflatten.h"
|
|
|
|
#include "archive/value_archive.h"
|
|
#include "core/operator.h"
|
|
|
|
namespace mmdeploy::graph {
|
|
|
|
void UnflattenNode::Build(TaskGraph& graph) {
|
|
auto p = graph.Add([this](Context& ctx) -> Result<void> {
|
|
auto args = std::move(ctx.current());
|
|
Value rets = Value::kArray;
|
|
auto idxs = from_value<std::vector<int>>(args.back());
|
|
for (int i = 0; i < rets.size() - 1; ++i) {
|
|
OUTCOME_TRY(auto ret, Unflatten(std::move(args[i]), idxs));
|
|
rets.push_back(std::move(ret));
|
|
}
|
|
ctx.current() = std::move(rets);
|
|
return success();
|
|
});
|
|
p->set_name(name());
|
|
}
|
|
|
|
class UnflattenCreator : public Creator<Node> {
|
|
public:
|
|
const char* GetName() const override { return "Unflatten"; }
|
|
int GetVersion() const override { return 0; }
|
|
std::unique_ptr<Node> Create(const Value& cfg) override {
|
|
return std::make_unique<UnflattenNode>(cfg);
|
|
}
|
|
};
|
|
|
|
REGISTER_MODULE(Node, UnflattenCreator);
|
|
|
|
} // namespace mmdeploy::graph
|