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

[WebNN EP] Support BatchNormalization op (#17071)

Browse source 
Adds support for BatchNormalization via WebNN meanVarianceNormalization.
This commit is contained in:
Wanming Lin 2023-08-16 08:52:09 +08:00 committed by GitHub
parent c0f8197157
commit 789bac1dc8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 72 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -134,6 +134,7 @@ static const InlinedHashMap<std::string, std::string> op_map = {
{"ArgMax", "argMax"},
{"ArgMin", "argMin"},
{"AveragePool", "averagePool2d"},
{"BatchNormalization", "meanVarianceNormalization"},
{"Cast", "cast"},
{"Ceil", "ceil"},
{"Clip", "clamp"},

118
onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc
View file

@ -34,7 +34,7 @@ Status NormalizationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder
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, "Layer/Instance/GroupNormalization requires at least two inputs.");
ORT_RETURN_IF_NOT(input_defs.size() >= 2, op_type, " requires at least two inputs.");
emscripten::val input = model_builder.GetOperand(input_defs[0]->Name());
std::vector<int64_t> input_shape;
@ -46,10 +46,27 @@ Status NormalizationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder
std::vector<int64_t> scale_shape;
ORT_RETURN_IF_NOT(GetShape(*input_defs[1], scale_shape, logger), "Cannot get scale shape");
const auto scale_size = scale_shape.size();
ORT_RETURN_IF_NOT(scale_size >= 1 && scale_size <= rank, "The scale size should be less than or equal to input size.");
// Except LayerNormalization, other normalization ops' scale input should be 1-D.
if (op_type == "LayerNormalization") {
ORT_RETURN_IF_NOT(scale_size >= 1 && scale_size <= rank,
"The scale size should be less than or equal to input size.");
} else {
ORT_RETURN_IF_NOT(scale_size == 1, "The scale size should be one.");
}
if (input_defs.size() >= 3 && !input_defs[2]->Name().empty()) {
// Bias input exists, and bias's shape should be the same as scale's shape.
std::vector<int64_t> bias_shape;
ORT_RETURN_IF_NOT(GetShape(*input_defs[2], bias_shape, logger), "Cannot get bias shape");
ORT_RETURN_IF_NOT(bias_shape == scale_shape, "The bias' shape should be equal to scale's shape.");
}
std::vector<uint32_t> new_scale_shape;
if (scale_size < rank) {
if (op_type == "LayerNormalization") {
if (op_type == "BatchNormalization") {
scale_shape.insert(scale_shape.begin(), 1);
scale_shape.insert(scale_shape.end(), rank - 2, 1);
} else if (op_type == "LayerNormalization") {
// Align right with leading ones.
scale_shape.insert(scale_shape.begin(), rank - scale_size, 1);
} else if (op_type == "InstanceNormalization") {
@ -66,48 +83,25 @@ Status NormalizationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder
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),
[](int64_t dim) -> int32_t { return SafeInt<int32_t>(dim); });
[](int64_t dim) -> uint32_t { return SafeInt<uint32_t>(dim); });
emscripten::val reshape_scale = model_builder.GetOperand(input_defs[1]->Name());
emscripten::val reshape_output_scale =
model_builder.GetBuilder().call<emscripten::val>("reshape", reshape_scale, emscripten::val::array(new_scale_shape));
options.set("scale", reshape_output_scale);
} else {
options.set("scale", model_builder.GetOperand(input_defs[1]->Name()));
}
if (input_defs.size() == 3) {
// Inputs contain optional bias
std::vector<int64_t> bias_shape;
ORT_RETURN_IF_NOT(GetShape(*input_defs[2], bias_shape, logger), "Cannot get bias shape");
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) {
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),
[](int64_t dim) -> int32_t { return SafeInt<int32_t>(dim); });
if (input_defs.size() >= 3 && !input_defs[2]->Name().empty()) {
// Bias input exists, and bias's shape is the same as scale's shape.
emscripten::val reshape_bias = model_builder.GetOperand(input_defs[2]->Name());
emscripten::val reshape_output_bias =
model_builder.GetBuilder().call<emscripten::val>("reshape", reshape_bias, emscripten::val::array(new_bias_shape));
model_builder.GetBuilder().call<emscripten::val>("reshape", reshape_bias, emscripten::val::array(new_scale_shape));
options.set("bias", reshape_output_bias);
} else {
}
} else {
options.set("scale", model_builder.GetOperand(input_defs[1]->Name()));
if (input_defs.size() >= 3 && !input_defs[2]->Name().empty()) {
// Bias input exists, and bias's shape is the same as scale's shape.
options.set("bias", model_builder.GetOperand(input_defs[2]->Name()));
}
}
@ -116,16 +110,35 @@ Status NormalizationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder
options.set("epsilon", helper.Get("epsilon", 1e-05f));
emscripten::val output = emscripten::val::undefined();
if (op_type == "LayerNormalization") {
if (op_type == "BatchNormalization") {
ORT_RETURN_IF_NOT(input_defs.size() == 5, "BatchNormalization requires five inputs.");
emscripten::val mean = model_builder.GetOperand(input_defs[3]->Name());
emscripten::val variance = model_builder.GetOperand(input_defs[4]->Name());
// Enlarge 1-D mean and variance to new scale shape.
emscripten::val reshape_mean =
model_builder.GetBuilder().call<emscripten::val>("reshape", mean, emscripten::val::array(new_scale_shape));
emscripten::val reshape_variance =
model_builder.GetBuilder().call<emscripten::val>("reshape", variance, emscripten::val::array(new_scale_shape));
std::vector<uint32_t> axes = {0};
for (uint32_t i = 2; i < rank; i++) {
axes.push_back(i);
}
options.set("axes", emscripten::val::array(axes));
options.set("mean", reshape_mean);
options.set("variance", reshape_variance);
output = model_builder.GetBuilder().call<emscripten::val>("meanVarianceNormalization", input, options);
} else if (op_type == "LayerNormalization") {
int64_t axis = helper.Get("axis", -1);
axis = HandleNegativeAxis(axis, rank);
std::vector<int32_t> axes(rank - narrow<size_t>(axis));
std::vector<uint32_t> axes(rank - SafeInt<uint32_t>(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++) {
std::vector<uint32_t> axes;
for (uint32_t i = 2; i < rank; i++) {
axes.emplace_back(i);
}
options.set("axes", emscripten::val::array(axes));
@ -133,21 +146,21 @@ Status NormalizationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder
} 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::vector<uint32_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); });
[](int64_t dim) -> uint32_t { return SafeInt<uint32_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));
new_shape.emplace_back(SafeInt<uint32_t>(input_shape[0]));
new_shape.emplace_back(SafeInt<uint32_t>(group_count));
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>())));
new_shape.emplace_back(SafeInt<uint32_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}));
options.set("axes", emscripten::val::array(std::vector<uint32_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));
@ -166,8 +179,11 @@ bool NormalizationOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initi
const WebnnDeviceType /* device_type */,
const logging::Logger& logger) const {
const auto& input_defs = node.InputDefs();
const auto& op_type = node.OpType();
NodeAttrHelper helper(node);
if (input_defs.size() < 2) {
LOGS(logger, VERBOSE) << "LayerNormalization requires at least two inputs.";
LOGS(logger, VERBOSE) << op_type << " requires at least two inputs.";
return false;
}
@ -179,7 +195,12 @@ bool NormalizationOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initi
const auto& output_defs = node.OutputDefs();
if (output_defs.size() != 1) {
LOGS(logger, VERBOSE) << node.OpType() << " output count must be one.";
LOGS(logger, VERBOSE) << op_type << " output count must be one.";
return false;
}
if (op_type == "BatchNormalization" && helper.Get("training_mode", 0)) {
LOGS(logger, VERBOSE) << "BatchNormalization with training_mode set to true is not supported.";
return false;
}
@ -192,6 +213,7 @@ void CreateNormalizationOpBuilder(const std::string& op_type, OpBuilderRegistrat
constexpr static std::string_view op_types[] =
{
"BatchNormalization",
"GroupNormalization",
"InstanceNormalization",
"LayerNormalization",

1
onnxruntime/core/providers/webnn/builders/op_builder_factory.cc
View file

@ -108,6 +108,7 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() {
}
{ // Normalization
CreateNormalizationOpBuilder("BatchNormalization", op_registrations);
CreateNormalizationOpBuilder("GroupNormalization", op_registrations);
CreateNormalizationOpBuilder("InstanceNormalization", op_registrations);
CreateNormalizationOpBuilder("LayerNormalization", op_registrations);