mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-19 19:00:47 +00:00
Add ML Program support for transpose op (#21364)
### Description Add support for transpose op ### Motivation and Context Enable support for Autodesk model
This commit is contained in:
parent
760a31c848
commit
fa287042ca
2 changed files with 23 additions and 6 deletions
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
#include "core/providers/coreml/builders/helper.h"
|
||||
#include "core/providers/coreml/builders/impl/base_op_builder.h"
|
||||
#include "core/providers/coreml/builders/impl/builder_utils.h"
|
||||
#include "core/providers/coreml/builders/model_builder.h"
|
||||
#include "core/providers/coreml/builders/op_builder_factory.h"
|
||||
#include "core/providers/coreml/shape_utils.h"
|
||||
|
|
@ -14,13 +15,13 @@ namespace coreml {
|
|||
class TransposeOpBuilder : public BaseOpBuilder {
|
||||
Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node,
|
||||
const logging::Logger& logger) const override;
|
||||
|
||||
bool SupportsMLProgram() const override { return true; }
|
||||
};
|
||||
|
||||
Status TransposeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
|
||||
const Node& node,
|
||||
const logging::Logger& logger) const {
|
||||
std::unique_ptr<COREML_SPEC::NeuralNetworkLayer> layer = model_builder.CreateNNLayer(node);
|
||||
|
||||
NodeAttrHelper helper(node);
|
||||
std::vector<int64_t> perm = helper.Get("perm", std::vector<int64_t>());
|
||||
std::vector<int64_t> input_shape;
|
||||
|
|
@ -33,12 +34,27 @@ Status TransposeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
|
|||
ORT_RETURN_IF_NOT(perm.size() == input_dims, "Perm and input should have same dimension");
|
||||
}
|
||||
|
||||
*layer->mutable_transpose()->mutable_axes() = {perm.cbegin(), perm.cend()};
|
||||
#if defined(COREML_ENABLE_MLPROGRAM)
|
||||
if (model_builder.CreateMLProgram()) {
|
||||
using namespace CoreML::Specification::MILSpec;
|
||||
|
||||
*layer->mutable_input()->Add() = node.InputDefs()[0]->Name();
|
||||
*layer->mutable_output()->Add() = node.OutputDefs()[0]->Name();
|
||||
std::unique_ptr<Operation> op = model_builder.CreateOperation(node, "transpose");
|
||||
AddOperationInput(*op, "x", node.InputDefs()[0]->Name());
|
||||
AddOperationInput(*op, "perm", model_builder.AddConstant(op->type(), "perm", perm));
|
||||
AddOperationOutput(*op, *node.OutputDefs()[0]);
|
||||
model_builder.AddOperation(std::move(op));
|
||||
|
||||
model_builder.AddLayer(std::move(layer));
|
||||
} else
|
||||
#endif // defined(COREML_ENABLE_MLPROGRAM)
|
||||
{
|
||||
std::unique_ptr<COREML_SPEC::NeuralNetworkLayer> layer = model_builder.CreateNNLayer(node);
|
||||
*layer->mutable_transpose()->mutable_axes() = {perm.cbegin(), perm.cend()};
|
||||
|
||||
*layer->mutable_input()->Add() = node.InputDefs()[0]->Name();
|
||||
*layer->mutable_output()->Add() = node.OutputDefs()[0]->Name();
|
||||
|
||||
model_builder.AddLayer(std::move(layer));
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -20,3 +20,4 @@ Keep in sync with doco generated from /docs/execution-providers/CoreML-Execution
|
|||
|ai.onnx:Sub||
|
||||
|ai.onnx:Sigmoid||
|
||||
|ai:onnx:Tanh||
|
||||
|ai:onnx:Transpose||
|
||||
|
|
|
|||
Loading…
Reference in a new issue