diff --git a/importerUtils.cpp b/importerUtils.cpp index ee6a56fb..89055701 100644 --- a/importerUtils.cpp +++ b/importerUtils.cpp @@ -1451,7 +1451,7 @@ nvinfer1::Dims makeDims(int32_t nbDims, int64_t val) } NodeOutputs normalizationHelper(ImporterContext* ctx, const ::ONNX_NAMESPACE::NodeProto& node, size_t const nodeIdx, - std::vector& inputs, bool const useV2) + std::vector& inputs, bool const useV2, bool const includeChannelAxis) { auto* input = &convertToTensor(inputs.at(0), ctx); auto* scale = &convertToTensor(inputs.at(1), ctx); @@ -1485,6 +1485,11 @@ NodeOutputs normalizationHelper(ImporterContext* ctx, const ::ONNX_NAMESPACE::No unsqueezeAxes.push_back(i); } + if (includeChannelAxis && nbGroups == 1) + { + axesMask |= 1u << 1; + } + scale = unsqueezeTensor(ctx, *scale, unsqueezeAxes); bias = unsqueezeTensor(ctx, *bias, unsqueezeAxes); diff --git a/importerUtils.hpp b/importerUtils.hpp index 31e2d67c..bfe69572 100644 --- a/importerUtils.hpp +++ b/importerUtils.hpp @@ -271,9 +271,9 @@ std::unique_ptr createPlugin(ImporterContext* ctx, ::ONNX_N // Helper function to return the identity of a TensorOrWeights TensorOrWeights identity(ImporterContext* ctx, TensorOrWeights input); -// Helper function to create normalization layers for GroupNorm and InstanceNorm +// Helper function to create normalization layers for GroupNorm and InstanceNorm. NodeOutputs normalizationHelper(ImporterContext* ctx, ::ONNX_NAMESPACE::NodeProto const& node, size_t const nodeIdx, - std::vector& inputs, bool const useV2); + std::vector& inputs, bool const useV2, bool const includeChannelAxis = false); // Given a list of axes in the range of [-rank, rank-1], where rank is the rank // of the corresponding data tensor, normalize to [0, rank-1]. diff --git a/onnxOpImporters.cpp b/onnxOpImporters.cpp index e8145975..4ffa69a1 100644 --- a/onnxOpImporters.cpp +++ b/onnxOpImporters.cpp @@ -2649,7 +2649,7 @@ DEFINE_BUILTIN_OP_IMPORTER(GreaterOrEqual) DEFINE_BUILTIN_OP_IMPORTER(GroupNormalization) { bool const useV2 = ctx->getOpsetVersion() >= 21; - return normalizationHelper(ctx, node, nodeIdx, inputs, useV2); + return normalizationHelper(ctx, node, nodeIdx, inputs, useV2, /*includeChannelAxis=*/true); } // singlePassShape is the shape of the output from a single pass. diff --git a/onnx_backend_test.py b/onnx_backend_test.py index f62004ad..7ecdfda7 100644 --- a/onnx_backend_test.py +++ b/onnx_backend_test.py @@ -158,6 +158,7 @@ backend_test.include(r'.*test_split.*custom.*') backend_test.include(r'.*test_instancenorm_.*_custom.*') backend_test.include(r'.*test_slice.*custom.*') +backend_test.include(r'.*test_group_normalization_.*') # exclude unenabled ops get pulled in with wildcards