mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
[WebNN EP] Support InstanceNormalization and GroupNormalization. (#16604)
### Description Add support for Op InstanceNormalization and GroupNormalization via MeanVarianceNormalization. ### Motivation and Context Enable more models like Olive'ified SD unet to run on WebNN EP.
This commit is contained in:
parent
15cb2f5a8a
commit
53057ec1f5
3 changed files with 91 additions and 33 deletions
|
|
@ -129,6 +129,8 @@ static const InlinedHashMap<std::string, std::string> op_map = {
|
|||
{"GlobalAveragePool", "averagePool2d"},
|
||||
{"GlobalMaxPool", "maxPool2d"},
|
||||
{"AveragePool", "averagePool2d"},
|
||||
{"GroupNormalization", "meanVarianceNormalization"},
|
||||
{"InstanceNormalization", "meanVarianceNormalization"},
|
||||
{"LayerNormalization", "meanVarianceNormalization"},
|
||||
{"MaxPool", "maxPool2d"},
|
||||
{"ReduceMax", "reduceMax"},
|
||||
|
|
|
|||
|
|
@ -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<int64_t> 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<int32_t> 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<int32_t> 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<int32_t> axes{static_cast<int32_t>(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<int32_t> axes(rank - axis);
|
||||
std::iota(axes.begin(), axes.end(), axis);
|
||||
options.set("axes", emscripten::val::array(axes));
|
||||
output = model_builder.GetBuilder().call<emscripten::val>("meanVarianceNormalization", input, options);
|
||||
} else if (op_type == "InstanceNormalization") {
|
||||
std::vector<int32_t> 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<emscripten::val>("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<int32_t> orig_shape, new_shape;
|
||||
std::transform(input_shape.cbegin(), input_shape.cend(),
|
||||
std::back_inserter(orig_shape),
|
||||
[](int64_t dim) -> int32_t { return SafeInt<int32_t>(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<int32_t>(input_shape[0]));
|
||||
new_shape.emplace_back(SafeInt<int32_t>(group_count));
|
||||
|
||||
emscripten::val output = model_builder.GetBuilder().call<emscripten::val>("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<int32_t>(std::reduce(input_shape.begin() + 2, input_shape.end(),
|
||||
input_shape[1] / group_count, std::multiplies<int64_t>())));
|
||||
// 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<int32_t>{2}));
|
||||
output = model_builder.GetBuilder().call<emscripten::val>("reshape", input, emscripten::val::array(new_shape));
|
||||
output = model_builder.GetBuilder().call<emscripten::val>("meanVarianceNormalization", output, options);
|
||||
output = model_builder.GetBuilder().call<emscripten::val>("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<NormalizationOpBuilder>());
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue