mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
minor fix in node unit change (#10405)
This commit is contained in:
parent
ea9c8a7cdc
commit
c6ef465011
2 changed files with 13 additions and 12 deletions
|
|
@ -857,10 +857,22 @@ void ReshapeOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const
|
|||
// onnxruntime::nnapi::Model.
|
||||
/* static */ bool ReshapeOpBuilder::CanSkipReshape(const ModelBuilder& model_builder, const NodeUnit& node_unit,
|
||||
size_t input_rank, size_t output_rank) {
|
||||
// Since we know this is a Reshape NodeUnit, so we can safely assume there is only 1 output
|
||||
// and the node_unit has only one output node.
|
||||
const auto& output_node_arg = node_unit.Outputs()[0].node_arg;
|
||||
const auto& output_name = output_node_arg.Name();
|
||||
const auto& output_node = *node_unit.GetOutputNodes()[0];
|
||||
|
||||
// Check if the Reshape output is a graph output, if so we cannot skip the Reshape
|
||||
// We do not care the case where the Reshape output is a dead end
|
||||
for (const auto* node_arg : model_builder.GetGraphViewer().GetOutputs()) {
|
||||
if (node_arg == &output_node_arg) {
|
||||
LOGS_DEFAULT(VERBOSE) << "Reshape/Flatten can not be skipped when the output is a graph output"
|
||||
<< ", output name, " << output_name;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// We will go through all the output edges
|
||||
for (auto it = output_node.OutputEdgesBegin(), end = output_node.OutputEdgesEnd(); it != end; ++it) {
|
||||
const auto& dest_node_unit = model_builder.GetNodeUnit(&it->GetNode());
|
||||
|
|
@ -897,17 +909,6 @@ void ReshapeOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const
|
|||
}
|
||||
}
|
||||
|
||||
// If we reach here, we have all the Reshape outputs are used by gemm/matmul, or Reshape has no output edge
|
||||
// Check if the Reshape output is a graph output, if so we cannot skip the Reshape
|
||||
// We do not care the case where the Reshape output is a dead end
|
||||
for (const auto* node_arg : model_builder.GetGraphViewer().GetOutputs()) {
|
||||
if (node_arg == &output_node_arg) {
|
||||
LOGS_DEFAULT(VERBOSE) << "Reshape/Flatten can not be skipped when the output is a graph output"
|
||||
<< ", output name, " << output_name;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
LOGS_DEFAULT(VERBOSE) << "Skipping Reshape/Flatten node ["
|
||||
<< node_unit.Name() << "] with output, " << output_name;
|
||||
return true;
|
||||
|
|
|
|||
|
|
@ -1236,7 +1236,7 @@ class ConcatOpSupportChecker : public BaseOpSupportChecker {
|
|||
bool ConcatOpSupportChecker::IsOpSupportedImpl(const InitializedTensorSet& /* initializers */, const NodeUnit& node_unit,
|
||||
const OpSupportCheckParams& /* params */) const {
|
||||
Shape input_shape;
|
||||
if (GetShape(node_unit.Inputs()[0].node_arg, input_shape))
|
||||
if (!GetShape(node_unit.Inputs()[0].node_arg, input_shape))
|
||||
return false;
|
||||
|
||||
const auto input_size = input_shape.size();
|
||||
|
|
|
|||
Loading…
Reference in a new issue