From e179f90638ab17fdda34ac298b10c7abe0baa837 Mon Sep 17 00:00:00 2001 From: Samaresh Kumar Singh Date: Wed, 20 May 2026 10:09:33 -0500 Subject: [PATCH] Fix GroupNormalization with num_groups=1 to include channel axis in reduction When num_groups equals 1 the ONNX importer was building axesMask from only the spatial dimensions, so INormalizationLayer ended up reducing per channel and the output silently matched InstanceNorm instead of the LayerNorm over CHW behavior that ONNX and PyTorch group_norm specify. This change threads an includeChannelAxis flag through normalizationHelper and adds axis 1 to axesMask only when num_groups is 1, leaving InstanceNormalization and the proper multi group case untouched. Fixes NVIDIA/TensorRT#4756 Signed-off-by: Samaresh Kumar Singh --- importerUtils.cpp | 7 ++++++- importerUtils.hpp | 4 ++-- onnxOpImporters.cpp | 2 +- onnx_backend_test.py | 1 + 4 files changed, 10 insertions(+), 4 deletions(-) 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