Struct dfdx::nn::BatchNorm2D

pub struct BatchNorm2D<const C: usize> {
    pub scale: Tensor1D<C>,
    pub bias: Tensor1D<C>,
    pub running_mean: Tensor1D<C>,
    pub running_var: Tensor1D<C>,
    pub epsilon: f32,
    pub momentum: f32,
}

Batch normalization for images as described in Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift

Generics:

• C the size of the spatial dimension to reduce. For 3d tensors this is the 0th dimension. For 4d tensors, this is the 1st dimension.

Training vs Inference

BatchNorm2D supports the following cases (see sections below for more details):

1. Training: ModuleMut and OwnedTape on the input tensor
2. Inference: Module and NoneTape on the input tensor.

NOTE: ModuleMut/NoneTape, and Module/OwnedTape will fail to compile.

Examples:

let bn: BatchNorm2D<3> = Default::default();
let _ = bn.forward(Tensor3D::<3, 2, 2>::zeros());
let _ = bn.forward(Tensor4D::<4, 3, 2, 2>::zeros());

Training

• Running statistics: updated with momentum
• Normalization: calculated using batch stats

Inference

• Running statistics: not updated
• Normalization: calculated using running stats

Fields

scale: Tensor1D<C>

Scale for affine transform. Defaults to 1.0

bias: Tensor1D<C>

Bias for affine transform. Defaults to 0.0

running_mean: Tensor1D<C>

Spatial mean that is updated during training. Defaults to 0.0

running_var: Tensor1D<C>

Spatial variance that is updated during training. Defaults to 1.0

epsilon: f32

Added to variance before taking sqrt for numerical stability. Defaults to 1e-5

momentum: f32

Controls exponential moving average of running stats.Defaults to 0.1

running_stat * (1.0 - momentum) + stat * momentum.

Trait Implementations

impl<const C: usize> CanUpdateWithGradients for BatchNorm2D<C>

fn update<G: GradientProvider>(
    &mut self,
    grads: &mut G,
    unused: &mut UnusedTensors
)

Updates self given the GradientProvider. When any parameters that are NOT present in G, then this function should add the tensor’s UniqueId to UnusedTensors.

impl<const C: usize> Clone for BatchNorm2D<C>

fn clone(&self) -> BatchNorm2D<C>

Returns a copy of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<const C: usize> Debug for BatchNorm2D<C>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<const C: usize> Default for BatchNorm2D<C>

fn default() -> Self

Returns the “default value” for a type.

impl<const C: usize> LoadFromNpz for BatchNorm2D<C>

fn read<R: Read + Seek>(
    &mut self,
    p: &str,
    r: &mut ZipArchive<R>
) -> Result<(), NpzError>

Reads this object from a ZipArchive. r with a base filename of filename_prefix.

fn load<P: AsRef<Path>>(&mut self, path: P) -> Result<(), NpzError>

Loads data from a .npz zip archive at the specified path.

impl<const C: usize, const H: usize, const W: usize> Module<Tensor3D<C, H, W, NoneTape>> for BatchNorm2D<C>

fn forward(&self, x: Tensor3D<C, H, W, NoneTape>) -> Self::Output

Inference 3d forward - does not update Self::running_mean and Self::running_var

type Output = Tensor3D<C, H, W, NoneTape>

The type that this unit produces given Input.

impl<const B: usize, const C: usize, const H: usize, const W: usize> Module<Tensor4D<B, C, H, W, NoneTape>> for BatchNorm2D<C>

fn forward(&self, x: Tensor4D<B, C, H, W, NoneTape>) -> Self::Output

Inference 4d forward - does not update Self::running_mean and Self::running_var

type Output = Tensor4D<B, C, H, W, NoneTape>

The type that this unit produces given Input.

impl<const C: usize, const H: usize, const W: usize> ModuleMut<Tensor3D<C, H, W, OwnedTape>> for BatchNorm2D<C>

fn forward_mut(&mut self, x: Tensor3D<C, H, W, OwnedTape>) -> Self::Output

Training 3d forward - updates Self::running_mean and Self::running_var

type Output = Tensor3D<C, H, W, OwnedTape>

The type that this unit produces given Input.

impl<const B: usize, const C: usize, const H: usize, const W: usize> ModuleMut<Tensor4D<B, C, H, W, OwnedTape>> for BatchNorm2D<C>

fn forward_mut(&mut self, x: Tensor4D<B, C, H, W, OwnedTape>) -> Self::Output

Training 4d forward - updates Self::running_mean and Self::running_var

type Output = Tensor4D<B, C, H, W, OwnedTape>

The type that this unit produces given Input.

impl<const C: usize> ResetParams for BatchNorm2D<C>

fn reset_params<R: Rng>(&mut self, _: &mut R)

Mutate the unit’s parameters using rand::Rng. Each implementor of this trait decides how the parameters are initialized. In fact, some impls may not even use the rng.

impl<const C: usize> SaveToNpz for BatchNorm2D<C>

fn write<W: Write + Seek>(&self, p: &str, w: &mut ZipWriter<W>) -> ZipResult<()>

Write this object into ZipWriter w with a base filename of filename_prefix.

fn save<P: AsRef<Path>>(&self, path: P) -> ZipResult<()>

Save this object into the .npz file determined located at path.