Implementation:LaurentMazare Tch rs Nn Batch Norm2d

Knowledge Sources	tch-rs
Domains	Deep_Learning, Optimization
Last Updated	2026-02-08 14:00 GMT

Overview

Concrete tool for applying batch normalization over 4D (N, C, H, W) inputs provided by the tch nn module.

Description

nn::batch_norm2d creates a BatchNorm layer with running_mean, running_var, and optional weight/bias tensors for the specified number of features. The layer implements ModuleT, computing batch normalization during training (using batch statistics) and inference (using running statistics) based on the train flag.

Usage

Use after convolutional layers in CNN architectures. Pass train: true during training and train: false during evaluation to switch between batch and running statistics.

Code Reference

Source Location

• Repository: tch-rs
• File: src/nn/batch_norm.rs
• Lines: 80-86

Signature

pub fn batch_norm2d<'a, T: Borrow<Path<'a>>>(
    vs: T,
    out_dim: i64,
    config: BatchNormConfig,
) -> BatchNorm

Import

use tch::nn;

I/O Contract

Inputs

Name	Type	Required	Description
vs	T: Borrow<Path>	Yes	VarStore path for parameter registration
out_dim	i64	Yes	Number of features (channels)
config	BatchNormConfig	Yes	Config: eps, momentum, affine (use Default::default())

Outputs

Name	Type	Description
BatchNorm	nn::BatchNorm	Layer with running_mean, running_var, optional weight/bias, implementing ModuleT

Usage Examples

use tch::{nn, nn::ModuleT, Device, Tensor};

let vs = nn::VarStore::new(Device::Cpu);
let bn = nn::batch_norm2d(vs.root() / "bn1", 32, Default::default());

let input = Tensor::randn([8, 32, 24, 24], tch::kind::FLOAT_CPU);
let output = bn.forward_t(&input, true);  // training mode

Related Pages

Implements Principle

• Principle:LaurentMazare_Tch_rs_Batch_Normalization