Struct dfdx::nn::modules::BatchNorm1D

pub struct BatchNorm1D<const C: usize, E: Dtype, D: DeviceStorage> {
    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: E,
    pub momentum: E,
}

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

Generics:

• C the size of the dimension to reduce. Both for 2d tensors (of the form <BATCH_SIZE, DIMENSION>) as well as 3d tensors (of the form <BATCH_SIZE, DIMENSION, SEQUENCE_LENGTH>), this is the 1st dimension.

Training vs Inference

BatchNorm1D 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 = BatchNorm1D<3>;
let bn = dev.build_module::<Model, f32>();
let _ = bn.forward(dev.zeros::<Rank2<4, 3>>());
let _ = bn.forward(dev.zeros::<Rank3<4, 3, 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: E

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

momentum: E

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 + DeviceStorage> Clone for BatchNorm1D<C, E, D>

fn clone(&self) -> BatchNorm1D<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 + DeviceStorage> Debug for BatchNorm1D<C, E, D>

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

Formats the value using the given formatter.

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

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

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

type Output = Tensor<(B, Const<C>), 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, L: Dim, E: Dtype, D: Device<E>> Module<Tensor<(B, Const<C>, L), E, D, NoneTape>> for BatchNorm1D<C, E, D>

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

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

type Output = Tensor<(B, Const<C>, L), 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, E: Dtype, D: Device<E>> ModuleMut<Tensor<(B, Const<C>), E, D, OwnedTape<E, D>>> for BatchNorm1D<C, E, D>

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

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

type Output = Tensor<(B, Const<C>), 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<B: Dim, const C: usize, L: Dim, E: Dtype, D: Device<E>> ModuleMut<Tensor<(B, Const<C>, L), E, D, OwnedTape<E, D>>> for BatchNorm1D<C, E, D>

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

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

type Output = Tensor<(B, Const<C>, L), 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 BatchNorm1D<C, E, D>

type To<E2: Dtype, D2: Device<E2>> = BatchNorm1D<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.