Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion importerUtils.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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<TensorOrWeights>& inputs, bool const useV2)
std::vector<TensorOrWeights>& inputs, bool const useV2, bool const includeChannelAxis)
{
auto* input = &convertToTensor(inputs.at(0), ctx);
auto* scale = &convertToTensor(inputs.at(1), ctx);
Expand Down Expand Up @@ -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);

Expand Down
4 changes: 2 additions & 2 deletions importerUtils.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -271,9 +271,9 @@ std::unique_ptr<nvinfer1::IPluginV3> 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<TensorOrWeights>& inputs, bool const useV2);
std::vector<TensorOrWeights>& 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].
Expand Down
2 changes: 1 addition & 1 deletion onnxOpImporters.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
1 change: 1 addition & 0 deletions onnx_backend_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down