class CV_EXPORTS GroupNormLayer : public Layer {

public:

static Ptr<GroupNormLayer> create(const LayerParams &params);

};

struct Layer_GroupNorm : public TestBaseWithParam<tuple<Backend, Target> >

{

void test_layer(const std::vector<int>& x_shape, int num_groups)

{

int backendId = get<0>(GetParam());

int targetId = get<1>(GetParam());

Mat x(x_shape, CV_32FC1);

Mat scale(x_shape[1], 1, CV_32FC1);

Mat b(x_shape[1], 1, CV_32FC1);

randu(x, 0.f, 1.f);

randu(scale, 0.f, 1.f);

randu(b, 0.f, 1.f);

Net net;

LayerParams lp;

lp.type = "GroupNormalization";

lp.name = "testLayer";

lp.set("num_groups", num_groups);

int id = net.addLayerToPrev(lp.name, lp.type, lp);

net.connect(0, 0, id, 0);

net.connect(0, 1, id, 1);

net.connect(0, 2, id, 2);

// warmup

{

std::vector<String> inpNames{"x", "scale", "b"};

net.setInputsNames(inpNames);

net.setInput(x, inpNames[0]);

net.setInput(scale, inpNames[1]);

net.setInput(b, inpNames[2]);

net.setPreferableBackend(backendId);

net.setPreferableTarget(targetId);

Mat out = net.forward();

}

TEST_CYCLE()

{

Mat res = net.forward();

}

SANITY_CHECK_NOTHING();

}

int N = 2;

int C = 64;

int H = 180;

int W = 240;

int num_groups = 16;

};

PERF_TEST_P_(Layer_GroupNorm, GroupNorm)

{

test_layer({N, C, H, W}, num_groups);

}

template <class T>

__global__ void normalize_mean_variance_groupwise(Span<T> output, View<T> input, View<T> scale, View<T> bias, View<float> means, View<float> inv_stddev, size_type inner_size, size_type C, size_type num_groups, size_type group_size) {

for (auto idx : grid_stride_range(output.size())) {

const index_type outer_idx = idx / inner_size;

const index_type c = outer_idx % C;

const index_type group_idx = outer_idx / group_size;

auto s = static_cast<float>(scale[c]) * inv_stddev[group_idx];

auto b = static_cast<float>(bias[c]);

output[idx] = (static_cast<float>(input[idx]) - means[group_idx]) * s + b;

}

}

template <class T>

void normalize_mean_variance_groupwise(const Stream& stream, Span<T> output, View<T> input, View<T> scale, View<T> bias, View<float> means, View<float> inv_stddev, std::size_t inner_size, std::size_t C, std::size_t num_groups, std::size_t group_size)

{

CV_Assert(input.size() == output.size());

CV_Assert(input.size() / inner_size == means.size() * group_size);

CV_Assert(means.size() == inv_stddev.size());

auto kernel = raw::normalize_mean_variance_groupwise<T>;

auto policy = make_policy(kernel, output.size(), 0, stream);

launch_kernel(kernel, policy, output, input, scale, bias, means, inv_stddev, inner_size, C, num_groups, group_size);

}

#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 530)

template void normalize_mean_variance_groupwise(const Stream&, Span<__half> /*output*/, View<__half> /*input*/, View<__half> /*scale*/, View<__half> /*bias*/, View<float> /*means*/, View<float> /*inv_stddev*/, std::size_t, std::size_t, std::size_t, std::size_t);

#endif

template void normalize_mean_variance_groupwise(const Stream&, Span<float> /*output*/, View<float> /*input*/, View<float> /*scale*/, View<float> /*bias*/, View<float> /*means*/, View<float> /*inv_stddev*/, std::size_t, std::size_t, std::size_t, std::size_t);

template <class T>

void normalize_mean_variance_groupwise(const csl::Stream &stream, csl::Span<T> output, csl::View<T> input, csl::View<T> scale, csl::View<T> bias, csl::View<float> means, csl::View<float> inv_stddev, std::size_t inner_size, std::size_t C, std::size_t num_groups, std::size_t group_size);

#ifndef OPENCV_DNN_SRC_CUDA4DNN_PRIMITIVES_GROUP_NORM_HPP

#define OPENCV_DNN_SRC_CUDA4DNN_PRIMITIVES_GROUP_NORM_HPP

#include "../../op_cuda.hpp"

#include "../csl/stream.hpp"

#include "../csl/span.hpp"

#include "../csl/tensor.hpp"

#include "../csl/workspace.hpp"

#include "../kernels/fill_copy.hpp"

#include "../kernels/mvn.hpp"

#include <opencv2/core.hpp>

#include <cstddef>

#include <vector>

#include <utility>

namespace cv { namespace dnn { namespace cuda4dnn {

template <class T>

class GroupNormOp final : public CUDABackendNode {

public:

using wrapper_type = GetCUDABackendWrapperType<T>;

GroupNormOp(csl::Stream stream_, float epsilon_, size_t loops, size_t num_groups)

: stream(std::move(stream_)), epsilon(epsilon_), num_groups(num_groups) {

csl::WorkspaceBuilder builder;

builder.require<float>(loops * num_groups); // mean and stdev for each group

builder.require<float>(loops * num_groups);

scratch_mem_in_bytes = builder.required_workspace_size();

}

void forward(const std::vector<cv::Ptr<BackendWrapper>>& inputs,

const std::vector<cv::Ptr<BackendWrapper>>& outputs,

csl::Workspace& workspace) override {

auto input_wrapper = inputs[0].dynamicCast<wrapper_type>();

auto scale_wrapper = inputs[1].dynamicCast<wrapper_type>();

auto bias_wrapper = inputs[2].dynamicCast<wrapper_type>();

auto input = input_wrapper->getView();

auto scale = scale_wrapper->getView();

auto bias = bias_wrapper->getView();

auto output_wrapper = outputs[0].dynamicCast<wrapper_type>();

auto output = output_wrapper->getSpan();

auto C = input.get_axis_size(1);

auto loops = input.size_range(0, 2);

auto norm_size = input.size_range(2, input.rank());

auto num_groups = this->num_groups;

auto group_size = C / num_groups;

if (norm_size == 1) {

kernels::fill<T>(stream, output, 0.f);

return;

} else {

auto ws_allocator = csl::WorkspaceAllocator(workspace);

auto mean = ws_allocator.get_span<float>(loops / group_size);

kernels::fill<float>(stream, mean, 0.f);

auto stdev = ws_allocator.get_span<float>(loops / group_size);

kernels::fill<float>(stream, stdev, 0.f);

kernels::reduce_mean_sqr_sum<T>(stream, mean, stdev, input, norm_size * group_size);

kernels::compute_normalization_scale(stream, stdev, mean, stdev, norm_size * group_size, epsilon);

kernels::normalize_mean_variance_groupwise<T>(stream, output, input, scale, bias, mean, stdev, norm_size, C, num_groups, group_size);

}

}

std::size_t get_workspace_memory_in_bytes() const noexcept override { return scratch_mem_in_bytes; }

private:

csl::Stream stream;

float epsilon;

std::size_t num_groups;

std::size_t scratch_mem_in_bytes;

};

}}} // cv::dnn::cuda4dnn

#endif // OPENCV_DNN_SRC_CUDA4DNN_PRIMITIVES_GROUP_NORM_HPP

CV_DNN_REGISTER_LAYER_CLASS(GroupNormalization, GroupNormLayer);

void fastNormGroup(const Mat &input, const Mat &scale, const Mat &bias, Mat &output, float epsilon, size_t num_groups) {

const auto input_shape = shape(input);

size_t N = input_shape[0], C = input_shape[1];

CV_CheckEQ(scale.total(), bias.total(), "fastNormGroup: scale and bias should have the same shape");

CV_CheckEQ(scale.total(), C, "fastNormGroup: scale should be a 1d tensor and match the channel of input");

CV_CheckGE(input.dims, 3, "fastNormGroup: input dimension >= 3");

size_t channels_per_group = C / num_groups;

size_t loops = N * num_groups;

size_t norm_size = static_cast<size_t>(total(input_shape, 2) * channels_per_group);

size_t step = norm_size / channels_per_group;

float inv_norm_size = 1.0 / norm_size;

auto fn = [&](const Range &r) {

const auto *input_data = input.ptr<const float>();

const auto *scale_data = scale.ptr<const float>();

const auto *bias_data = bias.ptr<const float>();

auto *output_data = output.ptr<float>();

for (int i = r.start; i < r.end; i++) {

const auto *x = input_data + norm_size * i;

auto *y = output_data + norm_size * i;

float mean = 0.f, mean_square = 0.f;

for (int j = 0; j < norm_size; j++) {

float v = x[j];

mean += v;

mean_square += v * v;

}

mean *= inv_norm_size;

mean_square = std::sqrt(std::max(0.f, mean_square * inv_norm_size - mean * mean) + epsilon);

float inv_stdev = 1.f / mean_square;

size_t group_idx = i % num_groups * channels_per_group;

for (size_t j = 0; j < norm_size; j++) {

size_t c = group_idx + (j / step);

float s = scale_data[c] * inv_stdev, b = bias_data[c];

y[j] = s * (x[j] - mean) + b;

}

}

};

double nstripes = loops * norm_size * (1 / 1024.0);

parallel_for_(Range(0, loops), fn, nstripes);

}

void fastNormGroup(const Mat &input, const Mat &scale, const Mat &bias, Mat &output, float epsilon, size_t num_groups);