saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 17:15:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[WebNN EP] Enable ArgMax and ArgMin for CPU backend (#20865)

Browse source 
WebNN TFLite backend supports ArgMax and ArgMin, but only supports
'select_last_index' value is 0.
This commit is contained in:
Wanming Lin 2024-06-04 05:12:11 +08:00 committed by GitHub
parent c128132dd8
commit 9c6481fa2d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 15 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
js/web/docs/webnn-operators.md
View file

@ -13,8 +13,8 @@ operators and the supported opset domain/versions in **WebNN EP** by ONNX Runtim
|:------:|:------:|:------:|:-:|:-:|:------|
| Abs | ai.onnx(7-12, 13+) | abs | ✓ | ✓ | |
| Add | ai.onnx(7-12, 13, 14+) | add | ✓ | ✓ | |
| ArgMax | ai.onnx(7-10, 11, 12, 13+) | argMax | ✗ | ✓ | |
| ArgMin | ai.onnx(7-10, 11, 12, 13+) | argMin | ✗ | ✓ | |
| ArgMax | ai.onnx(7-10, 11, 12, 13+) | argMax | ✓ | ✓ | WebNN CPU backend only supports 'select_last_index' value is 0 |
| ArgMin | ai.onnx(7-10, 11, 12, 13+) | argMin | ✓ | ✓ | WebNN CPU backend only supports 'select_last_index' value is 0 |
| AveragePool | ai.onnx(7-9, 10, 11, 12-18, 19+) | averagePool2d | ✓ | ✓ | Only supports 4-D input, 2-D 'kernel_shape', 'count_include_pad' value is 0 |
| BatchNormalization | ai.onnx(7-8, 9-13, 14, 15+) | batchNormalization | ✗ | ✓ | Only supports 'training_mode' value is 0, one output |
| Cast | ai.onnx(7-8, 9-12, 13-18, 19-20, 21+) | cast | ✗ | ✓ | |

4
onnxruntime/core/providers/webnn/builders/helper.h
View file

@ -157,8 +157,8 @@ std::vector<std::vector<NodeIndex>> GetSupportedNodes(const GraphViewer& graph_v
static const InlinedHashMap<std::string, WebnnOpInfo> op_map = {
{"Abs", {"abs", true}},
{"Add", {"add", true}},
{"ArgMax", {"argMax", false}},
{"ArgMin", {"argMin", false}},
{"ArgMax", {"argMax", true}},
{"ArgMin", {"argMin", true}},
{"AveragePool", {"averagePool2d", true}},
{"BatchNormalization", {"batchNormalization", false}},
{"Cast", {"cast", false}},

13
onnxruntime/core/providers/webnn/builders/impl/argmax_min_op_builder.cc
View file

@ -21,7 +21,7 @@ class ArgMaxMinOpBuilder : public BaseOpBuilder {
// Operator support related.
bool IsOpSupportedImpl(const InitializedTensorSet& initializers, const Node& node,
WebnnDeviceType /* device_type */, const logging::Logger& logger) const override;
WebnnDeviceType device_type, const logging::Logger& logger) const override;
bool HasSupportedInputsImpl(const Node& node, const WebnnDeviceType device_type,
const logging::Logger& logger) const override;
};
@ -68,7 +68,7 @@ Status ArgMaxMinOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
// Operator support related.
bool ArgMaxMinOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& /* initializers */,
const Node& node,
WebnnDeviceType /* device_type */,
WebnnDeviceType device_type,
const logging::Logger& logger) const {
const auto& input_defs = node.InputDefs();
@ -76,6 +76,15 @@ bool ArgMaxMinOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& /* initia
if (!GetShape(*input_defs[0], input_shape, logger))
return false;
// WebNN CPU backend only supports select_last_index = 0.
if (device_type == WebnnDeviceType::CPU) {
NodeAttrHelper helper(node);
const auto select_last_index = helper.Get("select_last_index", 0);
if (select_last_index) {
LOGS(logger, VERBOSE) << "ArgMax/ArgMin with select_last_index = 1 is not supported on WebNN CPU backend.";
return false;
}
}
return true;
}