Implementation:LaurentMazare Tch rs Group Norm

Knowledge Sources	LaurentMazare_Tch_rs
Domains	Neural Networks, Normalization, Deep Learning
Last Updated	2026-02-08 00:00 GMT

Overview

The group_norm module implements Group Normalization, which divides channels into groups and normalizes within each group, as described in the paper Group Normalization (Wu & He, 2018).

Description

GroupNorm normalizes input tensors by dividing the channel dimension into a specified number of groups and computing mean and variance within each group independently. This approach is less sensitive to batch size than Batch Normalization and does not require running statistics, making it well-suited for small-batch and per-instance scenarios.

The GroupNormConfig struct holds configuration parameters: cudnn_enabled (defaults to true), eps (epsilon for numerical stability, defaults to 1e-5), affine (whether to apply learnable scale and shift, defaults to true), ws_init (weight initialization, defaults to Const(1.)), and bs_init (bias initialization, defaults to Const(0.)).

The GroupNorm struct stores the config, optional weight and bias tensors (allocated only when affine is true), num_groups, and num_channels. The constructor function group_norm takes a variable store path, group count, channel count, and config. The layer implements Module::forward by delegating to Tensor::group_norm.

Usage

Use GroupNorm as a drop-in normalization layer in convolutional networks, especially when batch sizes are small or variable. It is commonly used in generative models, detection networks, and segmentation architectures.

Code Reference

Source Location

• Repository: LaurentMazare_Tch_rs
• File: src/nn/group_norm.rs

Signature

#[derive(Debug, Clone, Copy)]
pub struct GroupNormConfig {
    pub cudnn_enabled: bool,
    pub eps: f64,
    pub affine: bool,
    pub ws_init: super::Init,
    pub bs_init: super::Init,
}

#[derive(Debug)]
pub struct GroupNorm {
    config: GroupNormConfig,
    pub ws: Option<Tensor>,
    pub bs: Option<Tensor>,
    pub num_groups: i64,
    pub num_channels: i64,
}

pub fn group_norm<'a, T: Borrow<super::Path<'a>>>(
    vs: T,
    num_groups: i64,
    num_channels: i64,
    config: GroupNormConfig,
) -> GroupNorm;

Import

use tch::nn::{group_norm, GroupNormConfig};

I/O Contract

Parameter	Type	Description
vs	impl Borrow<Path>	Variable store path for parameter allocation
num_groups	i64	Number of groups to divide channels into
num_channels	i64	Total number of input channels (must be divisible by num_groups)
config	GroupNormConfig	Configuration struct

Config Field	Default Value	Description
cudnn_enabled	true	Enable cuDNN acceleration
eps	1e-5	Epsilon for numerical stability
affine	true	Learn scale (weight) and shift (bias) parameters
ws_init	Const(1.)	Weight initialization strategy
bs_init	Const(0.)	Bias initialization strategy

Forward Input	Forward Output
&Tensor of shape [N, C, ...] where C == num_channels	Tensor of same shape, group-normalized

Usage Examples

use tch::{nn, nn::Module, Device, Kind, Tensor};

let vs = nn::VarStore::new(Device::Cpu);
let root = vs.root();

// 8 groups over 32 channels
let gn = nn::group_norm(&root / "gn", 8, 32, Default::default());

// Forward pass
let input = Tensor::randn([4, 32, 16, 16], (Kind::Float, Device::Cpu));
let output = gn.forward(&input);
// output shape: [4, 32, 16, 16], normalized within each of the 8 groups

Related Pages

• Principle:LaurentMazare_Tch_rs_Group_Normalization