burn-nn 0.22.0

use crate::Initializer;
use burn_core as burn;

use burn::module::Flag;
use burn::module::{Content, DisplaySettings, ModuleDisplay};
use burn::tensor::{Device, Tensor, assert_shape};
use burn::{
    config::Config,
    module::{Module, Param, RunningState},
};

/// [`BatchNorm`] Configuration.
///
/// Used to create a [`BatchNorm`] layer using the [`BatchNormConfig::init`].
#[derive(Config, Debug)]
pub struct BatchNormConfig {
    /// The number of features. Must be greater than zero.
    pub num_features: usize,
    /// A finite, positive value required for numerical stability. Default: 1e-5
    #[config(default = 1e-5)]
    pub epsilon: f64,
    /// A finite momentum in `[0, 1]` used to update the running statistics. Default: 0.1
    ///
    /// A value of 0 preserves the running statistics; 1 replaces them with the current
    /// batch statistics.
    #[config(default = 0.1)]
    pub momentum: f64,
}

/// Applies Batch Normalization over a tensor.
///
/// Based upon the paper [Batch Normalization](https://arxiv.org/abs/1502.03167).
///
/// Assumes input tensor is of shape ``[batch_size, channels, ...]``.
///
/// `Y = norm(X) * γ + β`
///
/// Where:
/// - `X` is the input tensor
/// - `Y` is the output tensor
/// - `norm` is the normalization function
/// - `γ` is the learnable weight
/// - `β` is the learnable bias
///
/// Should be created using [`BatchNormConfig`].
#[derive(Module, Debug)]
#[module(custom_display)]
pub struct BatchNorm {
    /// The learnable weight gamma.
    pub gamma: Param<Tensor<1>>,
    /// The learnable weight beta.
    pub beta: Param<Tensor<1>>,
    /// Whether training behavior is enabled for this layer.
    pub training: Param<Flag>,
    /// The running mean.
    pub running_mean: RunningState<Tensor<1>>,
    /// The running variance.
    pub running_var: RunningState<Tensor<1>>,
    /// Momentum used to update the metrics.
    pub momentum: f64,
    /// A value required for numerical stability.
    pub epsilon: f64,
}

impl BatchNormConfig {
    /// Initializes a new [batch norm](BatchNorm) module.
    ///
    /// # Panics
    ///
    /// Panics if `num_features` is zero, `epsilon` is not finite and positive,
    /// or `momentum` is not finite and within `[0, 1]`.
    pub fn init(&self, device: &Device) -> BatchNorm {
        assert!(
            self.num_features > 0,
            "num_features must be greater than zero."
        );
        assert!(
            self.epsilon.is_finite() && self.epsilon > 0.0,
            "epsilon must be finite and positive."
        );
        assert!(
            self.momentum.is_finite() && (0.0..=1.0).contains(&self.momentum),
            "momentum must be finite and within [0, 1]."
        );

        let gamma = Initializer::Ones.init([self.num_features], device);
        let beta = Initializer::Zeros.init([self.num_features], device);

        let running_mean = Tensor::zeros([self.num_features], device);
        let running_var = Tensor::ones([self.num_features], device);

        BatchNorm {
            gamma,
            beta,
            training: Param::from_bool(true),
            running_mean: RunningState::new(running_mean),
            running_var: RunningState::new(running_var),
            momentum: self.momentum,
            epsilon: self.epsilon,
        }
    }
}

impl BatchNorm {
    /// Applies the forward pass on the input tensor.
    ///
    /// See [`BatchNorm`] for more information.
    ///
    /// # Shapes
    ///
    /// - `input`: ``[batch_size, num_features, ...]``
    /// - `output`: ``[batch_size, num_features, ...]``
    ///
    /// # Panics
    ///
    /// Panics if the input has rank < 2 or its second axis is not `num_features`.
    pub fn forward<const D: usize>(&self, input: Tensor<D>) -> Tensor<D> {
        let [num_features] = self.gamma.shape().dims();
        assert_shape!(input, [_, num_features, ..]);

        // Training behavior is selected by the device *and* the layer state. The device alone
        // says a backward is possible, not that this layer takes part in one: partial finetuning
        // freezes whole subtrees with [`freeze`](Module::freeze) or
        // [`freeze_group`](Module::freeze_group) and leaves them on the training device, because
        // that is where the rest of the graph lives.
        //
        // A frozen batch norm that still took the training path would recompute
        // the batch statistics — a second and third pass over an activation that
        // is already the layer's dominant cost — and then write them into
        // `running_mean` and `running_var`, mutating state the caller has said it does not want
        // trained. The identified flag lets a structural group select this behavior without
        // inferring the state of the whole layer from one tensor parameter.
        match input.device().is_autodiff() && self.training.is_enabled() {
            true => self.forward_train(input),
            false => self.forward_inference(input),
        }
    }

    fn forward_inference<const D: usize>(&self, input: Tensor<D>) -> Tensor<D> {
        let device = input.device();
        let mean = self.running_mean.value().to_device(&device);
        let var = self.running_var.value().to_device(&device);
        burn::tensor::module::batch_norm(
            input,
            self.gamma.val(),
            self.beta.val(),
            mean,
            var,
            self.epsilon,
        )
    }

    fn forward_train<const D: usize>(&self, input: Tensor<D>) -> Tensor<D> {
        let device = input.device();

        let result = burn::tensor::module::batch_norm_train(
            input,
            self.gamma.val(),
            self.beta.val(),
            self.epsilon,
        );

        let running_mean = self.running_mean.value_sync().to_device(&device);
        let running_var = self.running_var.value_sync().to_device(&device);

        let running_mean = running_mean
            .mul_scalar(1.0 - self.momentum)
            .add(result.mean.detach().mul_scalar(self.momentum));
        let running_var = running_var
            .mul_scalar(1.0 - self.momentum)
            .add(result.variance.detach().mul_scalar(self.momentum));

        self.running_mean.update(running_mean.detach());
        self.running_var.update(running_var.detach());

        result.output
    }
}

impl ModuleDisplay for BatchNorm {
    fn custom_settings(&self) -> Option<DisplaySettings> {
        DisplaySettings::new()
            .with_new_line_after_attribute(false)
            .optional()
    }

    fn custom_content(&self, content: Content) -> Option<Content> {
        let [num_features] = self.beta.shape().dims();

        let content = content
            .add("num_features", &num_features)
            .add("momentum", &self.momentum)
            .add("epsilon", &self.epsilon);
        match self.training.is_enabled() {
            true => content.optional(),
            false => content.add("training", &self.training).optional(),
        }
    }
}

#[cfg(feature = "std")]
#[cfg(test)]
mod tests_1d {
    use super::*;
    use burn::module::Module;
    use burn::tensor::TensorData;
    use burn::tensor::Tolerance;
    use rstest::rstest;
    type FT = f32;

    #[test]
    #[should_panic(expected = "num_features must be greater than zero.")]
    fn zero_features_is_rejected() {
        BatchNormConfig::new(0).init(&Default::default());
    }

    #[rstest]
    #[case::zero(0.0)]
    #[case::negative(-1e-5)]
    #[case::nan(f64::NAN)]
    #[case::positive_infinity(f64::INFINITY)]
    #[case::negative_infinity(f64::NEG_INFINITY)]
    #[should_panic(expected = "epsilon must be finite and positive.")]
    fn invalid_epsilon_is_rejected(#[case] epsilon: f64) {
        BatchNormConfig::new(1)
            .with_epsilon(epsilon)
            .init(&Default::default());
    }

    #[rstest]
    #[case::negative(-0.1)]
    #[case::above_one(1.1)]
    #[case::nan(f64::NAN)]
    #[case::positive_infinity(f64::INFINITY)]
    #[case::negative_infinity(f64::NEG_INFINITY)]
    #[should_panic(expected = "momentum must be finite and within [0, 1].")]
    fn invalid_momentum_is_rejected(#[case] momentum: f64) {
        BatchNormConfig::new(1)
            .with_momentum(momentum)
            .init(&Default::default());
    }

    #[rstest]
    #[case::zero(0.0, 0.0, 1.0)]
    #[case::one(1.0, 4.0, 4.0)]
    fn momentum_boundaries_update_running_statistics(
        #[case] momentum: f64,
        #[case] expected_mean: f32,
        #[case] expected_var: f32,
    ) {
        let device = Device::default().autodiff();
        let module = BatchNormConfig::new(1)
            .with_momentum(momentum)
            .init(&device);

        // Two batches verify that momentum 1 replaces the previous statistics.
        module.forward(Tensor::<2>::from_floats([[1.0], [3.0]], &device));
        let output = module.forward(Tensor::<2>::from_floats([[2.0], [6.0]], &device));

        // Both boundaries still use batch statistics to normalize during training.
        output
            .to_data()
            .assert_approx_eq::<FT>(&TensorData::from([[-1.0], [1.0]]), Tolerance::default());
        module
            .running_mean
            .value_sync()
            .to_data()
            .assert_approx_eq::<FT>(&TensorData::from([expected_mean]), Tolerance::default());
        module
            .running_var
            .value_sync()
            .to_data()
            .assert_approx_eq::<FT>(&TensorData::from([expected_var]), Tolerance::default());
    }

    #[test]
    #[should_panic(
        expected = "assert_shape!(input, [_, num_features, ..]): expected rank at least 2, got 1"
    )]
    fn input_rank_must_be_at_least_two() {
        let device = Default::default();
        let module = BatchNormConfig::new(3).init(&device);
        let _ = module.forward(Tensor::<1>::zeros([4], &device));
    }

    #[test]
    #[should_panic(
        expected = "assert_shape!(input, [_, num_features, ..]): axis 1 expected 3, got 4"
    )]
    fn input_channels_must_match() {
        let device = Default::default();
        let module = BatchNormConfig::new(3).init(&device);
        let _ = module.forward(Tensor::<3>::zeros([1, 4, 2], &device));
    }

    #[test]
    fn batch_norm_forward_train() {
        let device = Device::default().autodiff();
        let module = BatchNormConfig::new(3).init(&device);

        let output = module.forward(input_tensor(&device));

        output
            .to_data()
            .assert_approx_eq::<FT>(&expected_train(), Tolerance::rel_abs(0.1, 0.001));
    }

    #[test]
    fn batch_norm_forward_inference() {
        let device = Device::default();
        let device_autodiff = device.clone().autodiff();
        let module = BatchNormConfig::new(3).init(&device_autodiff);

        module.forward(input_tensor(&device_autodiff));
        let module = module.valid();
        let output = module.forward(input_tensor(&device));

        output
            .to_data()
            .assert_approx_eq::<FT>(&expected_valid(), Tolerance::default());
    }

    #[test]
    fn batch_norm_trains_under_gradient_checkpointing() {
        let device = Device::default().autodiff().gradient_checkpointing();
        let module = BatchNormConfig::new(3).init(&device);

        let output = module.forward(input_tensor(&device));

        output
            .to_data()
            .assert_approx_eq::<FT>(&expected_train(), Tolerance::rel_abs(0.1, 0.001));
        assert_eq!(module.running_mean.value_sync().dims(), [3]);
    }

    fn expected_valid() -> TensorData {
        TensorData::from([
            [[0.9409, 0.6976], [0.5892, 0.8774], [0.9106, 0.6844]],
            [[0.6012, 0.0782], [-0.0394, 0.9270], [0.6181, 0.5492]],
        ])
    }

    fn expected_train() -> TensorData {
        TensorData::from([
            [
                [1.1483e+00, 3.7521e-01],
                [1.6272e-03, 7.5067e-01],
                [1.6204e+00, -4.5168e-02],
            ],
            [
                [6.8856e-02, -1.5923e+00],
                [-1.6318e+00, 8.7949e-01],
                [-5.3368e-01, -1.0416e+00],
            ],
        ])
    }

    fn input_tensor(device: &Device) -> Tensor<3> {
        Tensor::<3>::from_floats(
            [
                [[0.9601, 0.7277], [0.6272, 0.9034], [0.9378, 0.7230]],
                [[0.6356, 0.1362], [0.0249, 0.9509], [0.6600, 0.5945]],
            ],
            device,
        )
    }

    #[test]
    fn batch_norm_forward_train_inference() {
        let device = Device::default();
        let device_autodiff = device.clone().autodiff();
        let module = BatchNormConfig::new(3).init(&device_autodiff);

        module.forward(input_tensor(&device_autodiff));
        let module = module.valid();
        let output = module.forward(input_tensor(&device));

        output
            .to_data()
            .assert_approx_eq::<FT>(&expected_valid(), Tolerance::default());

        let module = module.train();
        let output = module.forward(input_tensor(&device_autodiff));
        output
            .to_data()
            .assert_approx_eq::<FT>(&expected_train(), Tolerance::default());
    }
}

#[cfg(feature = "std")]
#[cfg(test)]
mod tests_2d {
    use super::*;
    use burn::module::Module;
    use burn::tensor::TensorData;
    use burn::tensor::Tolerance;
    type FT = f32;

    #[test]
    fn batch_norm_forward_train() {
        let device = Device::default().autodiff();
        let module = BatchNormConfig::new(3).init(&device);

        let output = module.forward(input_tensor(&device));

        let expected = TensorData::from([
            [
                [[1.5136, 0.7506], [-1.2216, 0.1477]],
                [[0.3135, 1.2252], [-0.4150, 0.6130]],
                [[1.4186, 0.3372], [-1.5183, 1.5262]],
            ],
            [
                [[0.4483, -1.1914], [-1.2010, 0.7537]],
                [[-1.6752, 1.3822], [-0.5058, -0.9381]],
                [[0.0200, -0.3097], [-0.5715, -0.9026]],
            ],
        ]);
        output
            .to_data()
            .assert_approx_eq::<FT>(&expected, Tolerance::rel_abs(0.1, 0.001));
    }

    #[test]
    fn batch_norm_forward_inference() {
        let device = Device::default();
        let device_autodiff = device.clone().autodiff();
        let module = BatchNormConfig::new(3).init(&device_autodiff);

        module.forward(input_tensor(&device_autodiff));
        let module = module.valid();
        let output = module.forward(input_tensor(&device));

        let expected = TensorData::from([
            [
                [[0.9538, 0.7103], [0.0808, 0.5179]],
                [[0.6015, 0.8910], [0.3703, 0.6966]],
                [[0.9171, 0.6912], [0.3037, 0.9395]],
            ],
            [
                [[0.6138, 0.0904], [0.0874, 0.7113]],
                [[-0.0297, 0.9408], [0.3415, 0.2042]],
                [[0.6250, 0.5561], [0.5013, 0.4323]],
            ],
        ]);
        output
            .to_data()
            .assert_approx_eq::<FT>(&expected, Tolerance::default());
    }

    #[test]
    fn batch_norm_running_mean() {
        let device = Device::default().autodiff();
        let module = BatchNormConfig::new(3).init(&device);

        let _output = module.forward(input_tensor(&device));

        let running_mean = module.running_mean.value_sync();

        let expected = TensorData::from([0.0499, 0.0532, 0.0656]);
        running_mean
            .reshape([3])
            .into_data()
            .assert_approx_eq::<FT>(&expected, Tolerance::default());
    }

    #[test]
    fn frozen_batch_norm_on_a_training_device_uses_the_running_statistics() {
        let device = Device::default().autodiff();
        // Frozen where partial finetuning leaves it: still on the training
        // device, because the rest of the graph is there, but not being trained.
        let module = BatchNormConfig::new(3).init(&device).freeze();

        let input = input_tensor(&device);
        let output = module.forward(input.clone());

        // Freshly initialized, the inference path is the identity: running mean
        // is zero, running variance is one, gamma is one and beta is zero. So
        // the input coming back out is proof the training path did not run —
        // that one normalizes the batch, and `batch_norm_forward_train` above
        // holds the quite different numbers it produces from this same input.
        output
            .to_data()
            .assert_approx_eq::<FT>(&input.to_data(), Tolerance::rel_abs(0.001, 0.001));
    }

    #[test]
    fn frozen_batch_norm_does_not_update_its_running_statistics() {
        let device = Device::default().autodiff();
        let module = BatchNormConfig::new(3).init(&device).freeze();

        let before = module.running_mean.value_sync().into_data();
        let _output = module.forward(input_tensor(&device));
        let after = module.running_mean.value_sync().into_data();

        // Freezing says the caller does not want this trained, and the running
        // statistics are state the training path writes. Untouched is the whole
        // point: a finetuning run that silently drifted them would corrupt the
        // frozen layer over its epochs and only show up at inference.
        after.assert_approx_eq::<FT>(&before, Tolerance::default());
    }

    #[test]
    fn enabling_gradients_does_not_make_running_statistics_trainable() {
        use burn::module::ParamGroup;

        let device = Device::default().autodiff();
        let module = BatchNormConfig::new(3).init(&device).freeze();
        let group = ParamGroup::ids_from_module(module.clone());

        let module = module.set_require_grad_group(group, true);

        assert!(module.gamma.is_require_grad());
        assert!(module.beta.is_require_grad());
        assert!(!module.training.is_enabled());
        assert!(!module.running_mean.value_sync().is_require_grad());
        assert!(!module.running_var.value_sync().is_require_grad());

        let grads = module.forward(input_tensor(&device)).sum().backward();

        assert!(module.gamma.grad(&grads).is_some());
        assert!(module.beta.grad(&grads).is_some());
        assert!(module.running_mean.value_sync().grad(&grads).is_none());
        assert!(module.running_var.value_sync().grad(&grads).is_none());
    }

    #[test]
    fn freezing_only_gamma_keeps_batch_norm_training_behavior() {
        use burn::module::ParamGroup;

        let device = Device::default().autodiff();
        let module = BatchNormConfig::new(3)
            .init(&device)
            .freeze_group(ParamGroup::from_path("gamma"));

        assert!(!module.gamma.is_require_grad());
        assert!(module.beta.is_require_grad());
        assert!(module.training.is_enabled());

        let before = module.running_mean.value_sync().into_data();
        let _output = module.forward(input_tensor(&device));
        let after = module.running_mean.value_sync().into_data();

        assert_ne!(before, after);
    }

    #[test]
    fn batch_norm_running_var() {
        let device = Device::default().autodiff();
        let module = BatchNormConfig::new(3).init(&device);

        let _output = module.forward(input_tensor(&device));

        let running_var = module.running_var.value_sync();

        let expected = TensorData::from([0.9106, 0.9105, 0.9045]);
        running_var
            .reshape([3])
            .into_data()
            .assert_approx_eq::<FT>(&expected, Tolerance::default());
    }

    #[test]
    fn batch_norm_running_mean_inner_module() {
        let device = Device::default().autodiff();
        let module = BatchNormConfig::new(3).init(&device);

        let _output = module.forward(input_tensor(&device));

        let module_valid = module.valid();
        let running_mean = module_valid.running_mean.value();
        let running_mean_after = module.running_mean.value();

        running_mean_after
            .into_data()
            .assert_approx_eq::<FT>(&running_mean.into_data(), Tolerance::default());
    }

    #[test]
    fn batch_norm_grads() {
        let device = Device::default().autodiff();
        let module = BatchNormConfig::new(3).init(&device);
        let input = input_tensor(&device).require_grad();

        let output = module.forward(input.clone());

        let grads = output.backward();

        let tolerance = Tolerance::rel_abs(0.1, 0.001);
        let expected = TensorData::from([0.0000e+00, -5.9035e-07, -6.0011e-07]);
        module
            .gamma
            .grad(&grads)
            .unwrap()
            .reshape([3])
            .into_data()
            .assert_approx_eq::<FT>(&expected, tolerance);

        let expected = TensorData::from([8., 8., 8.]);
        module
            .beta
            .grad(&grads)
            .unwrap()
            .reshape([3])
            .into_data()
            .assert_approx_eq::<FT>(&expected, tolerance);

        let expected = TensorData::from([
            [
                [[0.0000e+00, 0.0000e+00], [0.0000e+00, 0.0000e+00]],
                [[7.6400e-08, 2.9848e-07], [-1.0110e-07, 1.4933e-07]],
                [[5.3570e-07, 1.2732e-07], [-5.7336e-07, 5.7632e-07]],
            ],
            [
                [[0.0000e+00, 0.0000e+00], [0.0000e+00, 0.0000e+00]],
                [[-4.0807e-07, 3.3673e-07], [-1.2323e-07, -2.2854e-07]],
                [[7.5642e-09, -1.1695e-07], [-2.1582e-07, -3.4078e-07]],
            ],
        ]);
        input
            .grad(&grads)
            .unwrap()
            .into_data()
            .assert_approx_eq::<FT>(&expected, tolerance);
    }

    fn input_tensor(device: &Device) -> Tensor<4> {
        Tensor::<4>::from_floats(
            [
                [
                    [[0.9601, 0.7277], [0.1270, 0.5441]],
                    [[0.6272, 0.9034], [0.4066, 0.7179]],
                    [[0.9378, 0.7230], [0.3544, 0.9591]],
                ],
                [
                    [[0.6356, 0.1362], [0.1333, 0.7287]],
                    [[0.0249, 0.9509], [0.3791, 0.2481]],
                    [[0.6600, 0.5945], [0.5424, 0.4767]],
                ],
            ],
            device,
        )
    }

    #[test]
    fn display() {
        let batch_norm = BatchNormConfig::new(3).init(&Default::default());

        assert_eq!(
            alloc::format!("{batch_norm}"),
            "BatchNorm {num_features: 3, momentum: 0.1, epsilon: 0.00001, params: 12}"
        );

        let frozen = batch_norm.freeze();
        assert_eq!(
            alloc::format!("{frozen}"),
            "BatchNorm {num_features: 3, momentum: 0.1, epsilon: 0.00001, training: disabled, params: 12}"
        );
    }
}