Struct dfdx::nn::modules::BatchNorm2D

pub struct BatchNorm2D<const C: usize, E: Dtype, D: Storage<E>> {
    pub scale: Tensor<Rank1<C>, E, D>,
    pub bias: Tensor<Rank1<C>, E, D>,
    pub running_mean: Tensor<Rank1<C>, E, D>,
    pub running_var: Tensor<Rank1<C>, E, D>,
    pub epsilon: f64,
    pub momentum: f64,
}

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:

type Model = BatchNorm2D<3>;
let bn = dev.build_module::<Model, f32>();
let _ = bn.forward(dev.zeros::<Rank3<3, 2, 2>>());
let _ = bn.forward(dev.zeros::<Rank4<4, 3, 2, 2>>());

Training

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

Inference

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

Fields

scale: Tensor<Rank1<C>, E, D>

Scale for affine transform. Defaults to 1.0

bias: Tensor<Rank1<C>, E, D>

Bias for affine transform. Defaults to 0.0

running_mean: Tensor<Rank1<C>, E, D>

Spatial mean that is updated during training. Defaults to 0.0

running_var: Tensor<Rank1<C>, E, D>

Spatial variance that is updated during training. Defaults to 1.0

epsilon: f64

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

momentum: f64

Controls exponential moving average of running stats.Defaults to 0.1

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

Trait Implementations

impl<const C: usize, E: Clone + Dtype, D: Clone + Storage<E>> Clone for BatchNorm2D<C, E, D>

fn clone(&self) -> BatchNorm2D<C, E, D>

Returns a copy of the value.

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

Performs copy-assignment from source.

impl<const C: usize, E: Debug + Dtype, D: Debug + Storage<E>> Debug for BatchNorm2D<C, E, D>

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

Formats the value using the given formatter.

impl<B: Dim, const C: usize, H: Dim, W: Dim, E: Dtype, D: Device<E>> Module<Tensor<(B, Const<C>, H, W), E, D, NoneTape>> for BatchNorm2D<C, E, D>

fn try_forward( &self, x: Tensor<(B, Const<C>, H, W), E, D, NoneTape> ) -> Result<Self::Output, D::Err>

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

type Output = Tensor<(B, Const<C>, H, W), E, D, NoneTape>

The type that this unit produces given Input.

type Error = <D as HasErr>::Err

fn forward(&self, input: Input) -> Self::Output

Forward Input through the module and produce Module::Output.

impl<const C: usize, H: Dim, W: Dim, E: Dtype, D: Device<E>> Module<Tensor<(Const<C>, H, W), E, D, NoneTape>> for BatchNorm2D<C, E, D>

fn try_forward( &self, x: Tensor<(Const<C>, H, W), E, D, NoneTape> ) -> Result<Self::Output, D::Err>

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

type Output = Tensor<(Const<C>, H, W), E, D, NoneTape>

The type that this unit produces given Input.

type Error = <D as HasErr>::Err

fn forward(&self, input: Input) -> Self::Output

Forward Input through the module and produce Module::Output.

impl<B: Dim, const C: usize, H: Dim, W: Dim, E: Dtype, D: Device<E>> ModuleMut<Tensor<(B, Const<C>, H, W), E, D, OwnedTape<E, D>>> for BatchNorm2D<C, E, D>

fn try_forward_mut( &mut self, x: Tensor<(B, Const<C>, H, W), E, D, OwnedTape<E, D>> ) -> Result<Self::Output, D::Err>

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

type Output = Tensor<(B, Const<C>, H, W), E, D, OwnedTape<E, D>>

The type that this unit produces given Input.

type Error = <D as HasErr>::Err

fn forward_mut(&mut self, input: Input) -> Self::Output

Forward Input through the module and produce ModuleMut::Output.

impl<const C: usize, H: Dim, W: Dim, E: Dtype, D: Device<E>> ModuleMut<Tensor<(Const<C>, H, W), E, D, OwnedTape<E, D>>> for BatchNorm2D<C, E, D>

fn try_forward_mut( &mut self, x: Tensor<(Const<C>, H, W), E, D, OwnedTape<E, D>> ) -> Result<Self::Output, D::Err>

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

type Output = Tensor<(Const<C>, H, W), E, D, OwnedTape<E, D>>

The type that this unit produces given Input.

type Error = <D as HasErr>::Err

fn forward_mut(&mut self, input: Input) -> Self::Output

Forward Input through the module and produce ModuleMut::Output.

impl<const C: usize, E: Dtype, D: Device<E>> TensorCollection<E, D> for BatchNorm2D<C, E, D>

type To<E2: Dtype, D2: Device<E2>> = BatchNorm2D<C, E2, D2>

Type alias that specifies the how a module’s type changes when using a different dtype and/or device.

fn iter_tensors<V: ModuleVisitor<Self, E, D>>( visitor: &mut V ) -> Result<Option<Self::To<V::E2, V::D2>>, V::Err>

Specifies how to iterate through tensors or modules containted within this module, and how to contruct this module given values for its fields. Returns Err(_) to indicate an error, Ok(None) to indicate that there is no error and a module has not been built, and Ok(Some(_)) contains Self::Output<E2, D2>

fn module<F1, F2, Field>( name: &str, get_ref: F1, get_mut: F2 ) -> ModuleField<'_, F1, F2, Self, Field>where F1: FnMut(&Self) -> &Field, F2: FnMut(&mut Self) -> &mut Field, Field: TensorCollection<E, D>,

Creates a ModuleFields that represents a field that may contain one or more tensors.

fn tensor<F1, F2, S>( name: &str, get_ref: F1, get_mut: F2, options: TensorOptions<S, E, D> ) -> TensorField<'_, F1, F2, Self, S, E, D>where F1: FnMut(&Self) -> &Tensor<S, E, D>, F2: FnMut(&mut Self) -> &mut Tensor<S, E, D>, S: Shape,

Creates a ModuleFields that represents a tensor field.

fn scalar<F1, F2, N>( name: &str, get_ref: F1, get_mut: F2, options: ScalarOptions<N> ) -> ScalarField<'_, F1, F2, Self, N>where F1: FnMut(&Self) -> &N, F2: FnMut(&mut Self) -> &mut N, N: NumCast,

Creates a ModuleFields that represents a scalar field.