diff --git a/onnxruntime/core/providers/webnn/builders/helper.h b/onnxruntime/core/providers/webnn/builders/helper.h index 7b11f852b8..0b4b6e4687 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.h +++ b/onnxruntime/core/providers/webnn/builders/helper.h @@ -129,6 +129,8 @@ static const InlinedHashMap op_map = { {"GlobalAveragePool", "averagePool2d"}, {"GlobalMaxPool", "maxPool2d"}, {"AveragePool", "averagePool2d"}, + {"GroupNormalization", "meanVarianceNormalization"}, + {"InstanceNormalization", "meanVarianceNormalization"}, {"LayerNormalization", "meanVarianceNormalization"}, {"MaxPool", "maxPool2d"}, {"ReduceMax", "reduceMax"}, diff --git a/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc index bea4fe72d6..15c9b4acbe 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc @@ -27,11 +27,14 @@ class NormalizationOpBuilder : public BaseOpBuilder { const WebnnDeviceType /* device_type */, const logging::Logger& logger) const override; }; +// All normalization are based on layout NCHW. +// TODO: add support for NHWC. Status NormalizationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, const logging::Logger& logger) const { + const auto& op_type = node.OpType(); const auto& input_defs = node.InputDefs(); - ORT_RETURN_IF_NOT(input_defs.size() >= 2, "LayerNormalization requires at least two inputs."); + ORT_RETURN_IF_NOT(input_defs.size() >= 2, "Layer/Instance/GroupNormalization requires at least two inputs."); emscripten::val input = model_builder.GetOperand(input_defs[0]->Name()); std::vector input_shape; @@ -46,8 +49,23 @@ Status NormalizationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder ORT_RETURN_IF_NOT(scale_size >= 1 && scale_size <= rank, "The scale size should be less than or equal to input size."); if (scale_size < rank) { - // Enlarge new shape to input.rank, right aligned with leading ones - scale_shape.insert(scale_shape.begin(), rank - scale_size, 1); + if (op_type == "LayerNormalization") { + // Align right with leading ones. + scale_shape.insert(scale_shape.begin(), rank - scale_size, 1); + } else if (op_type == "InstanceNormalization") { + // Insert ones before and after the channel dimension. + scale_shape.insert(scale_shape.begin(), 1); + ORT_RETURN_IF(scale_size != 1 || rank < 2, + "The scale size should be 1 and rank should be at least 2 for InstanceNorm."); + scale_shape.insert(scale_shape.end(), rank - scale_size - 1, 1); + } else if (op_type == "GroupNormalization") { + // The input will be reshaped to 3D later. So just insert ones before the channel and after. + scale_shape.insert(scale_shape.begin(), 1); + scale_shape.insert(scale_shape.end(), 1); + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported normalization op: ", op_type); + } + std::vector new_scale_shape; std::transform(scale_shape.cbegin(), scale_shape.cend(), std::back_inserter(new_scale_shape), @@ -67,9 +85,20 @@ Status NormalizationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder const auto bias_size = bias_shape.size(); ORT_RETURN_IF_NOT(bias_size >= 1 && bias_size <= rank, "The bias size should be less than or equal to input size."); + // Enlarge new shape to input.rank. if (bias_size < rank) { - // Enlarge new shape to input.rank, right aligned with leading ones - bias_shape.insert(bias_shape.begin(), rank - bias_size, 1); + if (op_type == "LayerNormalization") { + bias_shape.insert(bias_shape.begin(), rank - bias_size, 1); + } else if (op_type == "InstanceNormalization") { + bias_shape.insert(bias_shape.begin(), 1); + bias_shape.insert(bias_shape.end(), rank - bias_size - 1, 1); + } else if (op_type == "GroupNormalization") { + bias_shape.insert(bias_shape.begin(), 1); + bias_shape.insert(bias_shape.end(), 1); + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported normalization op: ", op_type); + } + std::vector new_bias_shape; std::transform(bias_shape.cbegin(), bias_shape.cend(), std::back_inserter(new_bias_shape), @@ -86,12 +115,45 @@ Status NormalizationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder NodeAttrHelper helper(node); options.set("epsilon", helper.Get("epsilon", 1e-05f)); - int64_t axis = helper.Get("axis", -1); - axis = HandleNegativeAxis(axis, rank); - std::vector axes{static_cast(axis)}; - options.set("axes", emscripten::val::array(axes)); + emscripten::val output = emscripten::val::undefined(); + if (op_type == "LayerNormalization") { + int64_t axis = helper.Get("axis", -1); + axis = HandleNegativeAxis(axis, rank); + std::vector axes(rank - axis); + std::iota(axes.begin(), axes.end(), axis); + options.set("axes", emscripten::val::array(axes)); + output = model_builder.GetBuilder().call("meanVarianceNormalization", input, options); + } else if (op_type == "InstanceNormalization") { + std::vector axes; + for (size_t i = 2; i < rank; i++) { + axes.emplace_back(i); + } + options.set("axes", emscripten::val::array(axes)); + output = model_builder.GetBuilder().call("meanVarianceNormalization", input, options); + } else if (op_type == "GroupNormalization") { + ORT_RETURN_IF_NOT(helper.HasAttr("num_groups"), "GroupNormalization num_group must be provided."); + int32_t group_count = helper.Get("num_groups", -1); + std::vector orig_shape, new_shape; + std::transform(input_shape.cbegin(), input_shape.cend(), + std::back_inserter(orig_shape), + [](int64_t dim) -> int32_t { return SafeInt(dim); }); + // Add N and Group. + ORT_RETURN_IF_NOT(rank >= 2, "Input for GroupNormalization cannot be a scalar or 1D"); + new_shape.emplace_back(SafeInt(input_shape[0])); + new_shape.emplace_back(SafeInt(group_count)); - emscripten::val output = model_builder.GetBuilder().call("meanVarianceNormalization", input, options); + ORT_RETURN_IF_NOT(group_count > 0 && input_shape[1] % group_count == 0, + "GroupNormalization num_group must be divisible by group."); + new_shape.emplace_back(SafeInt(std::reduce(input_shape.begin() + 2, input_shape.end(), + input_shape[1] / group_count, std::multiplies()))); + // Input will be reshaped to (N, group count, channels per group x D1 x D2 ... Dn) and recovered after normalization. + options.set("axes", emscripten::val::array(std::vector{2})); + output = model_builder.GetBuilder().call("reshape", input, emscripten::val::array(new_shape)); + output = model_builder.GetBuilder().call("meanVarianceNormalization", output, options); + output = model_builder.GetBuilder().call("reshape", output, emscripten::val::array(orig_shape)); + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported normalization op: ", op_type); + } model_builder.AddOperand(node.OutputDefs()[0]->Name(), std::move(output)); return Status::OK(); @@ -114,30 +176,10 @@ bool NormalizationOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initi LOGS(logger, VERBOSE) << "Cannot get input shape."; return false; } - const auto rank = input_shape.size(); - - NodeAttrHelper helper(node); - int64_t axis = helper.Get("axis", -1); - axis = HandleNegativeAxis(axis, rank); - - const auto& scale_name = input_defs[1]->Name(); - if (!Contains(initializers, scale_name)) { - LOGS(logger, VERBOSE) << "The scale must be a constant initializer."; - return false; - } - - if (input_defs.size() == 3) { - // Inputs contain optional bias - const auto& bias_name = input_defs[2]->Name(); - if (!Contains(initializers, bias_name)) { - LOGS(logger, VERBOSE) << "The bias must be a constant initializer."; - return false; - } - } const auto& output_defs = node.OutputDefs(); if (output_defs.size() != 1) { - LOGS(logger, VERBOSE) << "MeanVarianceNormalization output count must be one."; + LOGS(logger, VERBOSE) << node.OpType() << " output count must be one."; return false; } @@ -145,8 +187,20 @@ bool NormalizationOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initi } void CreateNormalizationOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) { + if (op_registrations.op_builder_map.find(op_type) != op_registrations.op_builder_map.cend()) + return; + + constexpr static std::string_view op_types[] = + { + "GroupNormalization", + "InstanceNormalization", + "LayerNormalization", + }; + op_registrations.builders.push_back(std::make_unique()); - op_registrations.op_builder_map.emplace(op_type, op_registrations.builders.back().get()); + for (const auto& op_type : op_types) { + op_registrations.op_builder_map.emplace(op_type, op_registrations.builders.back().get()); + } } } // namespace webnn diff --git a/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc b/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc index 0a84fe07b8..b390853da1 100644 --- a/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc +++ b/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc @@ -90,7 +90,9 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() { CreateLogicalOpBuilder("Equal", op_registrations); } - { // LayerNormalization + { // Normalization + CreateNormalizationOpBuilder("GroupNormalization", op_registrations); + CreateNormalizationOpBuilder("InstanceNormalization", op_registrations); CreateNormalizationOpBuilder("LayerNormalization", op_registrations); }