diff --git a/onnxruntime/core/providers/webnn/builders/model_builder.cc b/onnxruntime/core/providers/webnn/builders/model_builder.cc index e8f116d390..097dd16307 100644 --- a/onnxruntime/core/providers/webnn/builders/model_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/model_builder.cc @@ -75,9 +75,18 @@ InitializedTensorSet ModelBuilder::GetInitializerTensors() { } void ModelBuilder::PreprocessInitializers() { + const auto& initializers = graph_viewer_.GetAllInitializedTensors(); const auto& node_indices = graph_viewer_.GetNodesInTopologicalOrder(); for (size_t i = 0; i < node_indices.size(); i++) { const auto* node(graph_viewer_.GetNode(node_indices[i])); + + // find all initializers consumed. AddInitializersToSkip will potentially decrement the usage count. + for (const auto* input : node->InputDefs()) { + if (input->Exists() && Contains(initializers, input->Name())) { + initializer_usage_[input->Name()]++; + } + } + if (const auto* op_builder = GetOpBuilder(*node)) { op_builder->AddInitializersToSkip(*this, *node); } @@ -90,12 +99,11 @@ Status ModelBuilder::RegisterInitializers() { const auto& name = tensor.name(); const auto& shape = tensor.dims(); - // Ignore the following tensors: - // 1. Empty tensors: optional tensors can be indicated by an empty name. - // 2. Tensors in skipped_initializers_: These are tensors that are not used as WebNN Constants. - // Note: Scalar tensors are excluded because ONNX Runtime will optimize same scalar initializers into one. - if (name.empty() || (Contains(skipped_initializers_, name) && !shape.empty())) + // skip initializer if there is no remaining usage + auto usage_count = initializer_usage_[name]; + if (usage_count == 0) { continue; + } std::vector dims; // When the shape is empty, it is scalar initializer that dims = {}; @@ -385,7 +393,13 @@ void ModelBuilder::AddOperand(const std::string& name, const emscripten::val& op } void ModelBuilder::AddInitializerToSkip(const std::string& tensor_name) { - skipped_initializers_.insert(tensor_name); + // Decrement usage count if this is a known initializer. + // For simplicity the OpBuilder::AddInitializersToSkip implementations may call this for arbitrary input names + // without first checking if the value is an initializer. + auto entry = initializer_usage_.find(tensor_name); + if (entry != initializer_usage_.end()) { + --entry->second; + } } void ModelBuilder::AddInputToSkip(const std::string& input_name) { diff --git a/onnxruntime/core/providers/webnn/builders/model_builder.h b/onnxruntime/core/providers/webnn/builders/model_builder.h index 0fc2fa2067..4e2d84f481 100644 --- a/onnxruntime/core/providers/webnn/builders/model_builder.h +++ b/onnxruntime/core/providers/webnn/builders/model_builder.h @@ -81,7 +81,7 @@ class ModelBuilder { InlinedHashMap input_output_info_; - InlinedHashSet skipped_initializers_; + std::unordered_map initializer_usage_; InlinedHashSet skipped_inputs_; uint32_t name_token_{0};