[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:
zesongw 2023-07-11 11:51:53 +08:00 committed by GitHub
parent 15cb2f5a8a
commit 53057ec1f5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 91 additions and 33 deletions

View file

@ -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"},

View file

@ -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

View file

@ -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);
}