class

GroupNormBackward

GroupNormBackward

source edit

Python wrappers

Autograd node for Group Normalization.

Computes per-(batch, group) mean and reciprocal standard deviation over the channels-in-group ( $C/G$ ) plus all spatial axes, then normalises and applies a per-channel (not per-group) affine $y = \gamma_c\,\hat{x}_{b,c,\mathbf{s}} + \beta_c$ . The saved statistics are required by the backward closed-form; spatial_dims_ retains the per-axis spatial sizes so the backward can recompute the reduction count.

Inherits FuncOp<GroupNormBackward, 3> with three saved-input slots — x, gamma, and beta.

Math

Let $g(c)$ denote the group containing channel $c$ (so $g(c) = \lfloor c\,G\,/\,C\rfloor$ ) and write $S = \prod_i \text{spatial\_dims}[i]$ . For each (b, g) pair:

\mu_{b,g} = \frac{1}{(C/G)\,S} \sum_{c\in g}\sum_{\mathbf{s}} x_{b,c,\mathbf{s}}, \qquad \sigma_{b,g}^2 = \frac{1}{(C/G)\,S} \sum_{c\in g}\sum_{\mathbf{s}} (x_{b,c,\mathbf{s}} - \mu_{b,g})^2

\hat{x}_{b,c,\mathbf{s}} = (x_{b,c,\mathbf{s}} - \mu_{b,g(c)}) \cdot \frac{1}{\sqrt{\sigma_{b,g(c)}^2 + \epsilon}}, \qquad y_{b,c,\mathbf{s}} = \gamma_c\,\hat{x}_{b,c,\mathbf{s}} + \beta_c

Mean and variance use the biased estimator.

References

Wu & He, "Group Normalization" (ECCV 2018, arXiv:1803.08494).

Attributes

schema_v1OpSchema

Registered schema (name="group_norm", version=1, amp_policy=ForceFP32, produces_grad=true).

saved_mean_Storage

Per-(batch, group) mean

\mu_{b,g}

, shape (B*G,).

saved_rstd_Storage

Per-(batch, group) reciprocal standard deviation

1 / \sqrt{\sigma_{b,g}^2 + \epsilon}

, shape (B*G,).

B_, C_, G_int

Batch, channel, and group counts captured from x.shape().

spatial_dims_std::vector<int>

Per-axis spatial sizes [S_0, S_1, ...] retained so the backward can recompute the spatial reduction count.

Notes

Thread safety: instances are created once during forward and consumed once during backward — no concurrent access expected.

Static methods

static

forward

→TensorImplPtr

int forward(const int & x, const int & gamma, const int & beta, int num_groups, double eps)

source edit

Run the forward pass.

Computes per-(batch, group) statistics, normalises x, applies the per-channel affine, caches saved_mean_ and saved_rstd_, and (when grad mode is on) registers this as the grad_fn of the returned output.

Parameters

xTensorImplPtr

Input tensor of shape (B, C, S_0, ...). Rank must be >= 2.

gammaTensorImplPtr

Per-channel scale of shape (C,).

betaTensorImplPtr

Per-channel shift of shape (C,).

num_groupsint

Number of groups G; must divide C evenly.

epsdouble

Numerical-stability constant added inside the square root.

Returns

TensorImplPtr

Normalised output, same shape as x.

Raises

ShapeMismatch

If C is not divisible by num_groups or if gamma and beta do not have shape (C,).

Methods

apply

→std::vector<Storage>

int apply(int grad_out)

source edit

Compute gradients for the three saved inputs.

Returns {dx, d_gamma, d_beta} in the saved-input order using the cached statistics. Closed-form summary (per channel $c$ ):

d\gamma_c = \sum_{b,\mathbf{s}} \hat{x}_{b,c,\mathbf{s}}\, \text{grad\_out}_{b,c,\mathbf{s}}, \qquad d\beta_c = \sum_{b,\mathbf{s}} \text{grad\_out}_{b,c,\mathbf{s}}

and the input gradient follows the standard GN backward formula restricted to each (b, g) group.

Parameters

grad_outStorage

Incoming gradient, same shape as the forward output.

Returns

std::vector<Storage>

Three-element vector {dx, d_gamma, d_beta}.

class

GroupNormBackward

GroupNormBackward

source edit

Python wrappers

Autograd node for Group Normalization.

Inherits FuncOp<GroupNormBackward, 3> with three saved-input slots — x, gamma, and beta.

Math

Let $g(c)$ denote the group containing channel $c$ (so $g(c) = \lfloor c\,G\,/\,C\rfloor$ ) and write $S = \prod_i \text{spatial\_dims}[i]$ . For each (b, g) pair:

\mu_{b,g} = \frac{1}{(C/G)\,S} \sum_{c\in g}\sum_{\mathbf{s}} x_{b,c,\mathbf{s}}, \qquad \sigma_{b,g}^2 = \frac{1}{(C/G)\,S} \sum_{c\in g}\sum_{\mathbf{s}} (x_{b,c,\mathbf{s}} - \mu_{b,g})^2

\hat{x}_{b,c,\mathbf{s}} = (x_{b,c,\mathbf{s}} - \mu_{b,g(c)}) \cdot \frac{1}{\sqrt{\sigma_{b,g(c)}^2 + \epsilon}}, \qquad y_{b,c,\mathbf{s}} = \gamma_c\,\hat{x}_{b,c,\mathbf{s}} + \beta_c

Mean and variance use the biased estimator.

References

Wu & He, "Group Normalization" (ECCV 2018, arXiv:1803.08494).

Attributes

schema_v1OpSchema

Registered schema (name="group_norm", version=1, amp_policy=ForceFP32, produces_grad=true).

saved_mean_Storage

Per-(batch, group) mean

\mu_{b,g}

, shape (B*G,).

saved_rstd_Storage

Per-(batch, group) reciprocal standard deviation

1 / \sqrt{\sigma_{b,g}^2 + \epsilon}

, shape (B*G,).

B_, C_, G_int

Batch, channel, and group counts captured from x.shape().

spatial_dims_std::vector<int>

Per-axis spatial sizes [S_0, S_1, ...] retained so the backward can recompute the spatial reduction count.

Notes

Thread safety: instances are created once during forward and consumed once during backward — no concurrent access expected.

Static methods

static

forward

→TensorImplPtr

int forward(const int & x, const int & gamma, const int & beta, int num_groups, double eps)

source edit

Run the forward pass.

Parameters

xTensorImplPtr

Input tensor of shape (B, C, S_0, ...). Rank must be >= 2.

gammaTensorImplPtr

Per-channel scale of shape (C,).

betaTensorImplPtr

Per-channel shift of shape (C,).

num_groupsint

Number of groups G; must divide C evenly.

epsdouble

Numerical-stability constant added inside the square root.

Returns

TensorImplPtr

Normalised output, same shape as x.

Raises

ShapeMismatch

If C is not divisible by num_groups or if gamma and beta do not have shape (C,).

Methods

apply

→std::vector<Storage>

int apply(int grad_out)

source edit

Compute gradients for the three saved inputs.

Returns {dx, d_gamma, d_beta} in the saved-input order using the cached statistics. Closed-form summary (per channel $c$ ):

d\gamma_c = \sum_{b,\mathbf{s}} \hat{x}_{b,c,\mathbf{s}}\, \text{grad\_out}_{b,c,\mathbf{s}}, \qquad d\beta_c = \sum_{b,\mathbf{s}} \text{grad\_out}_{b,c,\mathbf{s}}

and the input gradient follows the standard GN backward formula restricted to each (b, g) group.

Parameters

grad_outStorage

Incoming gradient, same shape as the forward output.

Returns

std::vector<Storage>

Three-element vector {dx, d_gamma, d_beta}.