Struct BatchNorm

pub struct BatchNorm<F: Float + Debug + Send + Sync> { /* private fields */ }

Batch Normalization layer

Implements batch normalization as described in “Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift” by Ioffe and Szegedy.

This normalization is applied along the feature dimension (channel dimension for CNNs) over a mini-batch of data.

Examples

use scirs2_neural::layers::{BatchNorm, Layer};
use ndarray::{Array, Array4};
use rand::rngs::SmallRng;
use rand::SeedableRng;

// Create a batch normalization layer for 3 channels
let mut rng = SmallRng::seed_from_u64(42);
let batch_norm = BatchNorm::new(3, 0.9, 1e-5, &mut rng).unwrap();

// Forward pass with a batch of 2 samples, 3 channels, 4x4 spatial dimensions
let batch_size = 2;
let channels = 3;
let height = 4;
let width = 4;
let input = Array4::<f64>::from_elem((batch_size, channels, height, width), 0.1).into_dyn();
let output = batch_norm.forward(&input).unwrap();

// Output shape should match input shape
assert_eq!(output.shape(), input.shape());

Implementations

impl<F: Float + Debug + ScalarOperand + Send + Sync + 'static> BatchNorm<F>

pub fn new<R: Rng>( num_features: usize, momentum: f64, eps: f64, _rng: &mut R, ) -> Result<Self>

Create a new batch normalization layer

Arguments

• num_features - Number of features/channels to normalize
• momentum - Momentum for running mean/variance updates (default: 0.9)
• eps - Small constant for numerical stability (default: 1e-5)
• rng - Random number generator for initialization

Returns

• A new batch normalization layer

Examples found in repository

examples/semantic_segmentation_complete.rs (line 194)

    pub fn new(in_channels: usize, out_channels: usize, rng: &mut SmallRng) -> StdResult<Self> {
        Ok(Self {
            conv1: Conv2D::new(
                in_channels,
                out_channels,
                (3, 3),
                (1, 1),
                PaddingMode::Same,
                rng,
            )?,
            bn1: BatchNorm::new(out_channels, 0.1, 1e-5, rng)?,
            conv2: Conv2D::new(
                out_channels,
                out_channels,
                (3, 3),
                (1, 1),
                PaddingMode::Same,
                rng,
            )?,
            bn2: BatchNorm::new(out_channels, 0.1, 1e-5, rng)?,
            pool: MaxPool2D::new((2, 2), (2, 2), None)?,
        })
    }

    pub fn forward(&self, input: &ArrayD<f32>) -> StdResult<(ArrayD<f32>, ArrayD<f32>)> {
        // First conv + bn + relu
        let x = self.conv1.forward(input)?;
        let x = self.bn1.forward(&x)?;
        // Note: In practice, add ReLU activation here

        // Second conv + bn + relu
        let x = self.conv2.forward(&x)?;
        let skip = self.bn2.forward(&x)?;
        // Note: In practice, add ReLU activation here

        // Pooling for downsampling
        let pooled = self.pool.forward(&skip)?;

        Ok((pooled, skip))
    }
}

/// U-Net style decoder block
pub struct DecoderBlock {
    conv1: Conv2D<f32>,
    bn1: BatchNorm<f32>,
    conv2: Conv2D<f32>,
    bn2: BatchNorm<f32>,
}

impl DecoderBlock {
    pub fn new(in_channels: usize, out_channels: usize, rng: &mut SmallRng) -> StdResult<Self> {
        Ok(Self {
            conv1: Conv2D::new(
                in_channels,
                out_channels,
                (3, 3),
                (1, 1),
                PaddingMode::Same,
                rng,
            )?,
            bn1: BatchNorm::new(out_channels, 0.1, 1e-5, rng)?,
            conv2: Conv2D::new(
                out_channels,
                out_channels,
                (3, 3),
                (1, 1),
                PaddingMode::Same,
                rng,
            )?,
            bn2: BatchNorm::new(out_channels, 0.1, 1e-5, rng)?,
        })
    }

    pub fn forward(
        &self,
        input: &ArrayD<f32>,
        skip: Option<&ArrayD<f32>>,
    ) -> StdResult<ArrayD<f32>> {
        // Upsample input (simplified - in practice use transpose convolution)
        let upsampled = self.upsample(input)?;

        // Concatenate with skip connection if provided
        let x = if let Some(skip_tensor) = skip {
            self.concatenate(&upsampled, skip_tensor)?
        } else {
            upsampled
        };

        // First conv + bn + relu
        let x = self.conv1.forward(&x)?;
        let x = self.bn1.forward(&x)?;
        // Note: In practice, add ReLU activation here

        // Second conv + bn + relu
        let x = self.conv2.forward(&x)?;
        let x = self.bn2.forward(&x)?;
        // Note: In practice, add ReLU activation here

        Ok(x)
    }

    fn upsample(&self, input: &ArrayD<f32>) -> StdResult<ArrayD<f32>> {
        // Simplified upsampling using nearest neighbor
        let shape = input.shape();
        let batch_size = shape[0];
        let channels = shape[1];
        let height = shape[2];
        let width = shape[3];

        let mut upsampled = Array4::<f32>::zeros((batch_size, channels, height * 2, width * 2));

        for b in 0..batch_size {
            for c in 0..channels {
                for i in 0..height {
                    for j in 0..width {
                        let value = input[[b, c, i, j]];
                        upsampled[[b, c, i * 2, j * 2]] = value;
                        upsampled[[b, c, i * 2, j * 2 + 1]] = value;
                        upsampled[[b, c, i * 2 + 1, j * 2]] = value;
                        upsampled[[b, c, i * 2 + 1, j * 2 + 1]] = value;
                    }
                }
            }
        }

        Ok(upsampled.into_dyn())
    }

    fn concatenate(&self, input1: &ArrayD<f32>, input2: &ArrayD<f32>) -> StdResult<ArrayD<f32>> {
        // Simplified concatenation along channel dimension
        let shape1 = input1.shape();
        let shape2 = input2.shape();

        if shape1[0] != shape2[0] || shape1[2] != shape2[2] || shape1[3] != shape2[3] {
            return Err("Shapes incompatible for concatenation".into());
        }

        let batch_size = shape1[0];
        let channels1 = shape1[1];
        let channels2 = shape2[1];
        let height = shape1[2];
        let width = shape1[3];

        let mut concatenated =
            Array4::<f32>::zeros((batch_size, channels1 + channels2, height, width));

        // Copy first tensor
        for b in 0..batch_size {
            for c in 0..channels1 {
                for i in 0..height {
                    for j in 0..width {
                        concatenated[[b, c, i, j]] = input1[[b, c, i, j]];
                    }
                }
            }
        }

        // Copy second tensor
        for b in 0..batch_size {
            for c in 0..channels2 {
                for i in 0..height {
                    for j in 0..width {
                        concatenated[[b, channels1 + c, i, j]] = input2[[b, c, i, j]];
                    }
                }
            }
        }

        Ok(concatenated.into_dyn())
    }
}

/// U-Net model for semantic segmentation
pub struct UNetModel {
    encoders: Vec<EncoderBlock>,
    decoders: Vec<DecoderBlock>,
    bottleneck: Sequential<f32>,
    final_conv: Conv2D<f32>,
    config: SegmentationConfig,
}

impl UNetModel {
    pub fn new(config: SegmentationConfig, rng: &mut SmallRng) -> StdResult<Self> {
        let mut encoders = Vec::new();
        let mut decoders = Vec::new();

        // Build encoder blocks
        let mut in_channels = 3; // RGB input
        for &out_channels in &config.encoder_channels {
            encoders.push(EncoderBlock::new(in_channels, out_channels, rng)?);
            in_channels = out_channels;
        }

        // Bottleneck
        let bottleneck_channels = config.encoder_channels.last().copied().unwrap_or(512);
        let mut bottleneck = Sequential::new();
        bottleneck.add(Conv2D::new(
            bottleneck_channels,
            bottleneck_channels * 2,
            (3, 3),
            (1, 1),
            PaddingMode::Same,
            rng,
        )?);
        bottleneck.add(BatchNorm::new(bottleneck_channels * 2, 0.1, 1e-5, rng)?);
        bottleneck.add(Conv2D::new(
            bottleneck_channels * 2,
            bottleneck_channels,
            (3, 3),
            (1, 1),
            PaddingMode::Same,
            rng,
        )?);
        bottleneck.add(BatchNorm::new(bottleneck_channels, 0.1, 1e-5, rng)?);

        // Build decoder blocks
        in_channels = bottleneck_channels;
        for (i, &out_channels) in config.decoder_channels.iter().enumerate() {
            let decoder_in_channels =
                if config.skip_connections && i < config.encoder_channels.len() {
                    // Skip connections come from corresponding encoder layer (in reverse order)
                    let encoder_idx = config.encoder_channels.len() - 1 - i;
                    in_channels + config.encoder_channels[encoder_idx]
                } else {
                    in_channels
                };
            decoders.push(DecoderBlock::new(decoder_in_channels, out_channels, rng)?);
            in_channels = out_channels;
        }

        // Final classification layer
        let final_channels = config.decoder_channels.last().copied().unwrap_or(32);
        let final_conv = Conv2D::new(
            final_channels,
            config.num_classes,
            (1, 1),
            (1, 1),
            PaddingMode::Same,
            rng,
        )?;

        Ok(Self {
            encoders,
            decoders,
            bottleneck,
            final_conv,
            config,
        })
    }

examples/model_architecture_visualization.rs (line 132)

fn create_cnn_model(rng: &mut SmallRng) -> Result<Sequential<f32>> {
    let mut model = Sequential::new();

    // First convolutional block
    let conv1 = Conv2D::new(
        1,                      // input channels
        32,                     // output channels
        (3, 3),                 // kernel size
        (1, 1),                 // stride
        PaddingMode::Custom(1), // padding mode
        rng,
    )?;
    model.add_layer(conv1);

    let batch_norm1 = BatchNorm::new(32, 0.99, 1e-5, rng)?;
    model.add_layer(batch_norm1);

    // Second convolutional block
    let conv2 = Conv2D::new(
        32,                     // input channels
        64,                     // output channels
        (3, 3),                 // kernel size
        (2, 2),                 // stride (downsampling)
        PaddingMode::Custom(1), // padding mode
        rng,
    )?;
    model.add_layer(conv2);

    let batch_norm2 = BatchNorm::new(64, 0.99, 1e-5, rng)?;
    model.add_layer(batch_norm2);

    let dropout1 = Dropout::new(0.25, rng)?;
    model.add_layer(dropout1);

    // Flatten for fully connected layers - commented out as Flatten is not available
    // let flatten = Flatten::new()?;
    // model.add_layer(flatten);

    // Dense layers
    let dense1 = Dense::new(64 * 14 * 14, 128, Some("relu"), rng)?;
    model.add_layer(dense1);

    let dropout2 = Dropout::new(0.5, rng)?;
    model.add_layer(dropout2);

    // Output layer
    let dense2 = Dense::new(128, 10, Some("softmax"), rng)?;
    model.add_layer(dense2);

    Ok(model)
}

examples/generative_models_complete.rs (line 169)

    pub fn new(config: GenerativeConfig, rng: &mut SmallRng) -> StdResult<Self> {
        let (_height, _width) = config.input_size;

        // Feature extraction layers
        let mut feature_extractor = Sequential::new();
        feature_extractor.add(Conv2D::new(1, 32, (3, 3), (2, 2), PaddingMode::Same, rng)?);
        feature_extractor.add(BatchNorm::new(32, 1e-5, 0.1, rng)?);

        feature_extractor.add(Conv2D::new(32, 64, (3, 3), (2, 2), PaddingMode::Same, rng)?);
        feature_extractor.add(BatchNorm::new(64, 1e-5, 0.1, rng)?);

        feature_extractor.add(Conv2D::new(
            64,
            128,
            (3, 3),
            (2, 2),
            PaddingMode::Same,
            rng,
        )?);
        feature_extractor.add(BatchNorm::new(128, 1e-5, 0.1, rng)?);

        feature_extractor.add(AdaptiveMaxPool2D::new((4, 4), None)?);

        // Calculate flattened feature size
        let feature_size = 128 * 4 * 4;

        // Mean head
        let mut mean_head = Sequential::new();
        mean_head.add(Dense::new(
            feature_size,
            config.hidden_dims[0],
            Some("relu"),
            rng,
        )?);
        mean_head.add(Dropout::new(0.2, rng)?);
        mean_head.add(Dense::new(
            config.hidden_dims[0],
            config.latent_dim,
            None,
            rng,
        )?);

        // Log variance head
        let mut logvar_head = Sequential::new();
        logvar_head.add(Dense::new(
            feature_size,
            config.hidden_dims[0],
            Some("relu"),
            rng,
        )?);
        logvar_head.add(Dropout::new(0.2, rng)?);
        logvar_head.add(Dense::new(
            config.hidden_dims[0],
            config.latent_dim,
            None,
            rng,
        )?);

        Ok(Self {
            feature_extractor,
            mean_head,
            logvar_head,
            config,
        })
    }

    pub fn forward(&self, input: &ArrayD<f32>) -> StdResult<(ArrayD<f32>, ArrayD<f32>)> {
        // Extract features
        let features = self.feature_extractor.forward(input)?;

        // Flatten features
        let batch_size = features.shape()[0];
        let feature_dim = features.len() / batch_size;
        let flattened = features
            .to_shape(IxDyn(&[batch_size, feature_dim]))?
            .to_owned();

        // Get mean and log variance
        let mean = self.mean_head.forward(&flattened)?;
        let logvar = self.logvar_head.forward(&flattened)?;

        Ok((mean, logvar))
    }
}

/// VAE Decoder that reconstructs images from latent codes
pub struct VAEDecoder {
    latent_projection: Sequential<f32>,
    feature_layers: Sequential<f32>,
    output_conv: Conv2D<f32>,
    config: GenerativeConfig,
}

impl VAEDecoder {
    pub fn new(config: GenerativeConfig, rng: &mut SmallRng) -> StdResult<Self> {
        // Project latent to feature space
        let mut latent_projection = Sequential::new();
        latent_projection.add(Dense::new(
            config.latent_dim,
            config.hidden_dims[0],
            Some("relu"),
            rng,
        )?);
        latent_projection.add(Dense::new(
            config.hidden_dims[0],
            128 * 4 * 4,
            Some("relu"),
            rng,
        )?);

        // Feature reconstruction layers (simplified transpose convolutions)
        let mut feature_layers = Sequential::new();
        feature_layers.add(Conv2D::new(
            128,
            64,
            (3, 3),
            (1, 1),
            PaddingMode::Same,
            rng,
        )?);
        feature_layers.add(BatchNorm::new(64, 1e-5, 0.1, rng)?);

        feature_layers.add(Conv2D::new(64, 32, (3, 3), (1, 1), PaddingMode::Same, rng)?);
        feature_layers.add(BatchNorm::new(32, 1e-5, 0.1, rng)?);

        // Output layer
        let output_conv = Conv2D::new(32, 1, (3, 3), (1, 1), PaddingMode::Same, rng)?;

        Ok(Self {
            latent_projection,
            feature_layers,
            output_conv,
            config,
        })
    }

    pub fn forward(&self, latent: &ArrayD<f32>) -> StdResult<ArrayD<f32>> {
        // Project latent to feature space
        let projected = self.latent_projection.forward(latent)?;

        // Reshape to spatial format
        let batch_size = projected.shape()[0];
        let reshaped = projected.into_shape_with_order(IxDyn(&[batch_size, 128, 4, 4]))?;

        // Upsample to target size (simplified)
        let upsampled = self.upsample(&reshaped)?;

        // Apply feature layers
        let features = self.feature_layers.forward(&upsampled)?;

        // Generate output
        let output = self.output_conv.forward(&features)?;

        Ok(output)
    }

    fn upsample(&self, input: &ArrayD<f32>) -> StdResult<ArrayD<f32>> {
        let shape = input.shape();
        let batch_size = shape[0];
        let channels = shape[1];
        let height = shape[2];
        let width = shape[3];

        let (target_height, target_width) = self.config.input_size;
        let scale_h = target_height / height;
        let scale_w = target_width / width;

        let mut upsampled =
            Array4::<f32>::zeros((batch_size, channels, target_height, target_width));

        for b in 0..batch_size {
            for c in 0..channels {
                for i in 0..height {
                    for j in 0..width {
                        let value = input[[b, c, i, j]];
                        for di in 0..scale_h {
                            for dj in 0..scale_w {
                                let new_i = i * scale_h + di;
                                let new_j = j * scale_w + dj;
                                if new_i < target_height && new_j < target_width {
                                    upsampled[[b, c, new_i, new_j]] = value;
                                }
                            }
                        }
                    }
                }
            }
        }

        Ok(upsampled.into_dyn())
    }
}

/// Complete VAE model
pub struct VAEModel {
    encoder: VAEEncoder,
    decoder: VAEDecoder,
    config: GenerativeConfig,
}

impl VAEModel {
    pub fn new(config: GenerativeConfig, rng: &mut SmallRng) -> StdResult<Self> {
        let encoder = VAEEncoder::new(config.clone(), rng)?;
        let decoder = VAEDecoder::new(config.clone(), rng)?;

        Ok(Self {
            encoder,
            decoder,
            config,
        })
    }

    pub fn forward(
        &self,
        input: &ArrayD<f32>,
    ) -> StdResult<(ArrayD<f32>, ArrayD<f32>, ArrayD<f32>)> {
        // Encode
        let (mean, logvar) = self.encoder.forward(input)?;

        // Reparameterization trick (simplified)
        let latent = self.reparameterize(&mean, &logvar)?;

        // Decode
        let reconstruction = self.decoder.forward(&latent)?;

        Ok((reconstruction, mean, logvar))
    }

    fn reparameterize(&self, mean: &ArrayD<f32>, logvar: &ArrayD<f32>) -> StdResult<ArrayD<f32>> {
        // Sample epsilon from standard normal
        let mut epsilon = Array::zeros(mean.raw_dim());
        let mut rng = SmallRng::seed_from_u64(42); // Fixed seed for reproducibility

        for elem in epsilon.iter_mut() {
            *elem = rng.random_range(-1.0..1.0); // Approximate normal
        }

        // z = mean + std * epsilon, where std = exp(0.5 * logvar)
        let mut result = Array::zeros(mean.raw_dim());
        for (((&m, &lv), &eps), res) in mean
            .iter()
            .zip(logvar.iter())
            .zip(epsilon.iter())
            .zip(result.iter_mut())
        {
            let std = (0.5 * lv).exp();
            *res = m + std * eps;
        }

        Ok(result)
    }

    /// Generate new samples from random latent codes
    pub fn generate(&self, batch_size: usize) -> StdResult<ArrayD<f32>> {
        // Sample random latent codes
        let mut latent = Array2::<f32>::zeros((batch_size, self.config.latent_dim));
        let mut rng = SmallRng::seed_from_u64(123);

        for elem in latent.iter_mut() {
            *elem = rng.random_range(-1.0..1.0);
        }

        let latent_dyn = latent.into_dyn();

        // Decode to generate images
        self.decoder.forward(&latent_dyn)
    }

    /// Interpolate between two latent codes
    pub fn interpolate(
        &self,
        latent1: &ArrayD<f32>,
        latent2: &ArrayD<f32>,
        steps: usize,
    ) -> StdResult<Vec<ArrayD<f32>>> {
        let mut results = Vec::new();

        for i in 0..steps {
            let alpha = i as f32 / (steps - 1) as f32;

            // Linear interpolation
            let mut interpolated = Array::zeros(latent1.raw_dim());
            for ((&l1, &l2), interp) in latent1
                .iter()
                .zip(latent2.iter())
                .zip(interpolated.iter_mut())
            {
                *interp = (1.0 - alpha) * l1 + alpha * l2;
            }

            let generated = self.decoder.forward(&interpolated)?;
            results.push(generated);
        }

        Ok(results)
    }
}

/// VAE Loss combining reconstruction and KL divergence
pub struct VAELoss {
    reconstruction_loss: MeanSquaredError,
    beta: f32,
}

impl VAELoss {
    pub fn new(beta: f32) -> Self {
        Self {
            reconstruction_loss: MeanSquaredError::new(),
            beta,
        }
    }

    pub fn compute_loss(
        &self,
        reconstruction: &ArrayD<f32>,
        target: &ArrayD<f32>,
        mean: &ArrayD<f32>,
        logvar: &ArrayD<f32>,
    ) -> StdResult<(f32, f32, f32)> {
        // Reconstruction loss
        let recon_loss = self.reconstruction_loss.forward(reconstruction, target)?;

        // KL divergence loss: -0.5 * sum(1 + logvar - mean^2 - exp(logvar))
        let mut kl_loss = 0.0f32;
        for (&m, &lv) in mean.iter().zip(logvar.iter()) {
            kl_loss += -0.5 * (1.0 + lv - m * m - lv.exp());
        }
        kl_loss /= mean.len() as f32; // Average over elements

        let total_loss = recon_loss + self.beta * kl_loss;

        Ok((total_loss, recon_loss, kl_loss))
    }
}

/// Simple GAN Generator
pub struct GANGenerator {
    layers: Sequential<f32>,
    config: GenerativeConfig,
}

impl GANGenerator {
    pub fn new(config: GenerativeConfig, rng: &mut SmallRng) -> StdResult<Self> {
        let mut layers = Sequential::new();

        // Project noise to feature space
        layers.add(Dense::new(
            config.latent_dim,
            config.hidden_dims[0],
            Some("relu"),
            rng,
        )?);
        layers.add(BatchNorm::new(config.hidden_dims[0], 1e-5, 0.1, rng)?);

        layers.add(Dense::new(
            config.hidden_dims[0],
            config.hidden_dims[1] * 2,
            Some("relu"),
            rng,
        )?);
        layers.add(BatchNorm::new(config.hidden_dims[1] * 2, 1e-5, 0.1, rng)?);

        // Output layer
        let output_size = config.input_size.0 * config.input_size.1;
        layers.add(Dense::new(
            config.hidden_dims[1] * 2,
            output_size,
            Some("tanh"),
            rng,
        )?);

        Ok(Self { layers, config })
    }

examples/object_detection_complete.rs (line 184)

    pub fn new(config: DetectionConfig, rng: &mut SmallRng) -> StdResult<Self> {
        // Feature extraction backbone (simplified ResNet-like)
        let mut feature_extractor = Sequential::new();

        // Initial conv block
        feature_extractor.add(Conv2D::new(3, 64, (7, 7), (2, 2), PaddingMode::Same, rng)?);
        feature_extractor.add(BatchNorm::new(64, 0.1, 1e-5, rng)?);
        feature_extractor.add(MaxPool2D::new((2, 2), (2, 2), None)?);

        // Feature blocks
        feature_extractor.add(Conv2D::new(
            64,
            128,
            (3, 3),
            (2, 2),
            PaddingMode::Same,
            rng,
        )?);
        feature_extractor.add(BatchNorm::new(128, 0.1, 1e-5, rng)?);

        feature_extractor.add(Conv2D::new(
            128,
            256,
            (3, 3),
            (2, 2),
            PaddingMode::Same,
            rng,
        )?);
        feature_extractor.add(BatchNorm::new(256, 0.1, 1e-5, rng)?);

        // Global pooling to fixed size
        feature_extractor.add(AdaptiveMaxPool2D::new(config.feature_map_size, None)?);

        // Classification head
        let mut classifier_head = Sequential::new();
        let feature_dim = 256 * config.feature_map_size.0 * config.feature_map_size.1;
        classifier_head.add(Dense::new(feature_dim, 512, Some("relu"), rng)?);
        classifier_head.add(Dropout::new(0.5, rng)?);
        classifier_head.add(Dense::new(512, 256, Some("relu"), rng)?);
        classifier_head.add(Dropout::new(0.3, rng)?);
        classifier_head.add(Dense::new(
            256,
            config.num_classes * config.max_objects,
            Some("softmax"),
            rng,
        )?);

        // Bounding box regression head
        let mut bbox_regressor = Sequential::new();
        bbox_regressor.add(Dense::new(feature_dim, 512, Some("relu"), rng)?);
        bbox_regressor.add(Dropout::new(0.5, rng)?);
        bbox_regressor.add(Dense::new(512, 256, Some("relu"), rng)?);
        bbox_regressor.add(Dropout::new(0.3, rng)?);
        bbox_regressor.add(Dense::new(256, 4 * config.max_objects, None, rng)?); // 4 coordinates per object

        Ok(Self {
            feature_extractor,
            classifier_head,
            bbox_regressor,
            config,
        })
    }

examples/batchnorm_example.rs (line 25)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Batch Normalization Example");

    // Initialize random number generator with a fixed seed for reproducibility
    let mut rng = SmallRng::seed_from_u64(42);

    // Create a sample CNN architecture with batch normalization

    // 1. Define input dimensions
    let batch_size = 2;
    let in_channels = 3; // RGB input
    let height = 32;
    let width = 32;

    // 2. Create convolutional layer
    let conv = Conv2D::new(in_channels, 16, (3, 3), (1, 1), PaddingMode::Same, &mut rng)?;

    // 3. Create batch normalization layer for the conv output
    let batch_norm = BatchNorm::new(16, 0.9, 1e-5, &mut rng)?;

    // 4. Create random input data
    let input = Array4::<f32>::from_elem((batch_size, in_channels, height, width), 0.0);

    // Randomly fill input with values between -1 and 1
    let mut input_mut = input.clone();
    for n in 0..batch_size {
        for c in 0..in_channels {
            for h in 0..height {
                for w in 0..width {
                    input_mut[[n, c, h, w]] = rng.random_range(-1.0..1.0);
                }
            }
        }
    }

    // 5. Forward pass through conv layer
    println!("Input shape: {:?}", input_mut.shape());
    let conv_output = conv.forward(&input_mut.into_dyn())?;
    println!("Conv output shape: {:?}", conv_output.shape());

    // 6. Forward pass through batch normalization
    let bn_output = batch_norm.forward(&conv_output)?;
    println!("BatchNorm output shape: {:?}", bn_output.shape());

    // Print statistics of the conv output and batch norm output
    let conv_mean = compute_mean(&conv_output);
    let conv_std = compute_std(&conv_output, conv_mean);

    let bn_mean = compute_mean(&bn_output);
    let bn_std = compute_std(&bn_output, bn_mean);

    println!("\nStatistics before BatchNorm:");
    println!("  Mean: {:.6}", conv_mean);
    println!("  Std:  {:.6}", conv_std);

    println!("\nStatistics after BatchNorm:");
    println!("  Mean: {:.6}", bn_mean);
    println!("  Std:  {:.6}", bn_std);

    // Switch to inference mode
    let mut bn_inference = BatchNorm::new(16, 0.9, 1e-5, &mut rng)?;

    // First do a forward pass in training mode to accumulate statistics
    bn_inference.forward(&conv_output)?;

    // Now switch to inference mode
    bn_inference.set_training(false);
    let bn_inference_output = bn_inference.forward(&conv_output)?;

    let bn_inference_mean = compute_mean(&bn_inference_output);
    let bn_inference_std = compute_std(&bn_inference_output, bn_inference_mean);

    println!("\nStatistics in inference mode:");
    println!("  Mean: {:.6}", bn_inference_mean);
    println!("  Std:  {:.6}", bn_inference_std);

    // Example of using BatchNorm in a simple neural network
    println!("\nExample: BatchNorm in a simple neural network");

    // Create a random 2D input (batch_size, features)
    let batch_size = 16;
    let in_features = 10;
    let mut input_2d = Array::from_elem((batch_size, in_features), 0.0);

    // Randomly fill input with values between -1 and 1
    for n in 0..batch_size {
        for f in 0..in_features {
            input_2d[[n, f]] = rng.random_range(-1.0..1.0);
        }
    }

    // Create dense layer
    let dense1 = Dense::new(in_features, 32, None, &mut rng)?;

    // Create batch norm for dense output
    let bn1 = BatchNorm::new(32, 0.9, 1e-5, &mut rng)?;

    // Forward passes
    let dense1_output = dense1.forward(&input_2d.into_dyn())?;
    let bn1_output = bn1.forward(&dense1_output)?;

    println!(
        "Dense output stats - Mean: {:.6}, Std: {:.6}",
        compute_mean(&dense1_output),
        compute_std(&dense1_output, compute_mean(&dense1_output))
    );

    println!(
        "After BatchNorm - Mean: {:.6}, Std: {:.6}",
        compute_mean(&bn1_output),
        compute_std(&bn1_output, compute_mean(&bn1_output))
    );

    Ok(())
}

examples/dropout_example.rs (line 99)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Dropout Example");

    // Initialize random number generator with a fixed seed for reproducibility
    let mut rng = SmallRng::seed_from_u64(42);

    // 1. Example: Simple dropout on a vector
    println!("\nExample 1: Simple dropout on a vector");

    // Create a vector of ones
    let n_features = 20;
    let input = Array2::<f32>::from_elem((1, n_features), 1.0);

    // Create a dropout layer with dropout probability 0.5
    let dropout = Dropout::new(0.5, &mut rng)?;

    // Apply dropout in training mode
    let output = dropout.forward(&input.into_dyn())?;

    // Count elements that were dropped
    let mut dropped_count = 0;
    for i in 0..n_features {
        let val = output.slice(ndarray::s![0, i]).into_scalar();
        if *val == 0.0 {
            dropped_count += 1;
        }
    }

    // Print results
    println!("Input: Array of {} ones", n_features);
    println!(
        "Output after dropout (p=0.5): {} elements dropped",
        dropped_count
    );
    println!(
        "Dropout rate: {:.2}%",
        (dropped_count as f32 / n_features as f32) * 100.0
    );

    // 2. Example: Using dropout in a neural network to prevent overfitting
    println!("\nExample 2: Using dropout in a simple neural network");

    // Create a simple neural network with dropout
    // 1. Create dense layer (input -> hidden)
    let input_dim = 10;
    let hidden_dim = 100;
    let output_dim = 2;
    let dense1 = Dense::new(input_dim, hidden_dim, Some("relu"), &mut rng)?;

    // 2. Create dropout layer (after first dense layer)
    let dropout1 = Dropout::new(0.2, &mut rng)?;

    // 3. Create second dense layer (hidden -> output)
    let dense2 = Dense::new(hidden_dim, output_dim, None, &mut rng)?;

    // Generate a random input
    let batch_size = 16;
    let mut input_data = Array2::<f32>::zeros((batch_size, input_dim));
    for i in 0..batch_size {
        for j in 0..input_dim {
            input_data[[i, j]] = rng.random_range(-1.0..1.0);
        }
    }

    // Forward pass through the network with dropout
    println!("Running forward pass with dropout...");
    let hidden_output = dense1.forward(&input_data.clone().into_dyn())?;
    let hidden_dropped = dropout1.forward(&hidden_output)?;
    let final_output = dense2.forward(&hidden_dropped)?;

    println!("Input shape: {:?}", input_data.shape());
    println!("Hidden layer shape: {:?}", hidden_output.shape());
    println!("Output shape: {:?}", final_output.shape());

    // 3. Example: Dropout in a CNN architecture
    println!("\nExample 3: Dropout in a CNN architecture");

    // Create a sample CNN architecture with dropout

    // 1. Create convolutional layer
    let in_channels = 3; // RGB input
    let out_channels = 16;
    let conv1 = Conv2D::new(
        in_channels,
        out_channels,
        (3, 3),
        (1, 1),
        PaddingMode::Same,
        &mut rng,
    )?;

    // 2. Create batch normalization layer
    let bn1 = BatchNorm::new(out_channels, 0.9, 1e-5, &mut rng)?;

    // 3. Create dropout layer for spatial dropout (dropping entire feature maps)
    let dropout_conv = Dropout::new(0.25, &mut rng)?;

    // 4. Create flattened dense layer
    let height = 32;
    let width = 32;
    let flattened_size = out_channels * height * width;
    let dense3 = Dense::new(flattened_size, output_dim, None, &mut rng)?;

    // Generate random input image
    let mut image_input = Array4::<f32>::zeros((batch_size, in_channels, height, width));
    for n in 0..batch_size {
        for c in 0..in_channels {
            for h in 0..height {
                for w in 0..width {
                    image_input[[n, c, h, w]] = rng.random_range(-1.0..1.0);
                }
            }
        }
    }

    // Forward pass through CNN with dropout
    println!("Running forward pass through CNN with dropout...");
    let conv_output = conv1.forward(&image_input.clone().into_dyn())?;
    let bn_output = bn1.forward(&conv_output)?;
    let dropout_output = dropout_conv.forward(&bn_output)?;

    // Reshape for dense layer (flatten spatial and channel dimensions)
    let mut flattened = Array2::<f32>::zeros((batch_size, flattened_size));
    for n in 0..batch_size {
        let mut idx = 0;
        for c in 0..out_channels {
            for h in 0..height {
                for w in 0..width {
                    flattened[[n, idx]] =
                        *dropout_output.slice(ndarray::s![n, c, h, w]).into_scalar();
                    idx += 1;
                }
            }
        }
    }

    let cnn_output = dense3.forward(&flattened.into_dyn())?;

    println!("CNN output shape: {:?}", cnn_output.shape());

    // 4. Example: Switching between training and inference modes
    println!("\nExample 4: Switching between training and inference modes");

    // Create dropout layer with p=0.5
    let mut switchable_dropout = Dropout::new(0.5, &mut rng)?;

    // Create some input data
    let test_input = Array2::<f32>::from_elem((1, 10), 1.0);

    // Training mode (default)
    let training_output = switchable_dropout.forward(&test_input.clone().into_dyn())?;

    // Count non-zero elements in training mode
    let training_nonzero = training_output.iter().filter(|&&x| x != 0.0).count();

    // Switch to inference mode
    switchable_dropout.set_training(false);
    let inference_output = switchable_dropout.forward(&test_input.clone().into_dyn())?;

    // Count non-zero elements in inference mode
    let inference_nonzero = inference_output.iter().filter(|&&x| x != 0.0).count();

    println!("Training mode: {} of 10 elements kept", training_nonzero);
    println!("Inference mode: {} of 10 elements kept", inference_nonzero);

    // In inference mode, all elements should be preserved
    assert_eq!(inference_nonzero, 10);

    println!("\nDropout implementation demonstration completed successfully!");

    Ok(())
}

pub fn set_training(&mut self, training: bool)

Set the training mode

In training mode, batch statistics are used for normalization and running statistics are updated. In inference mode, running statistics are used for normalization and not updated.

Examples found in repository

examples/batchnorm_example.rs (line 73)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Batch Normalization Example");

    // Initialize random number generator with a fixed seed for reproducibility
    let mut rng = SmallRng::seed_from_u64(42);

    // Create a sample CNN architecture with batch normalization

    // 1. Define input dimensions
    let batch_size = 2;
    let in_channels = 3; // RGB input
    let height = 32;
    let width = 32;

    // 2. Create convolutional layer
    let conv = Conv2D::new(in_channels, 16, (3, 3), (1, 1), PaddingMode::Same, &mut rng)?;

    // 3. Create batch normalization layer for the conv output
    let batch_norm = BatchNorm::new(16, 0.9, 1e-5, &mut rng)?;

    // 4. Create random input data
    let input = Array4::<f32>::from_elem((batch_size, in_channels, height, width), 0.0);

    // Randomly fill input with values between -1 and 1
    let mut input_mut = input.clone();
    for n in 0..batch_size {
        for c in 0..in_channels {
            for h in 0..height {
                for w in 0..width {
                    input_mut[[n, c, h, w]] = rng.random_range(-1.0..1.0);
                }
            }
        }
    }

    // 5. Forward pass through conv layer
    println!("Input shape: {:?}", input_mut.shape());
    let conv_output = conv.forward(&input_mut.into_dyn())?;
    println!("Conv output shape: {:?}", conv_output.shape());

    // 6. Forward pass through batch normalization
    let bn_output = batch_norm.forward(&conv_output)?;
    println!("BatchNorm output shape: {:?}", bn_output.shape());

    // Print statistics of the conv output and batch norm output
    let conv_mean = compute_mean(&conv_output);
    let conv_std = compute_std(&conv_output, conv_mean);

    let bn_mean = compute_mean(&bn_output);
    let bn_std = compute_std(&bn_output, bn_mean);

    println!("\nStatistics before BatchNorm:");
    println!("  Mean: {:.6}", conv_mean);
    println!("  Std:  {:.6}", conv_std);

    println!("\nStatistics after BatchNorm:");
    println!("  Mean: {:.6}", bn_mean);
    println!("  Std:  {:.6}", bn_std);

    // Switch to inference mode
    let mut bn_inference = BatchNorm::new(16, 0.9, 1e-5, &mut rng)?;

    // First do a forward pass in training mode to accumulate statistics
    bn_inference.forward(&conv_output)?;

    // Now switch to inference mode
    bn_inference.set_training(false);
    let bn_inference_output = bn_inference.forward(&conv_output)?;

    let bn_inference_mean = compute_mean(&bn_inference_output);
    let bn_inference_std = compute_std(&bn_inference_output, bn_inference_mean);

    println!("\nStatistics in inference mode:");
    println!("  Mean: {:.6}", bn_inference_mean);
    println!("  Std:  {:.6}", bn_inference_std);

    // Example of using BatchNorm in a simple neural network
    println!("\nExample: BatchNorm in a simple neural network");

    // Create a random 2D input (batch_size, features)
    let batch_size = 16;
    let in_features = 10;
    let mut input_2d = Array::from_elem((batch_size, in_features), 0.0);

    // Randomly fill input with values between -1 and 1
    for n in 0..batch_size {
        for f in 0..in_features {
            input_2d[[n, f]] = rng.random_range(-1.0..1.0);
        }
    }

    // Create dense layer
    let dense1 = Dense::new(in_features, 32, None, &mut rng)?;

    // Create batch norm for dense output
    let bn1 = BatchNorm::new(32, 0.9, 1e-5, &mut rng)?;

    // Forward passes
    let dense1_output = dense1.forward(&input_2d.into_dyn())?;
    let bn1_output = bn1.forward(&dense1_output)?;

    println!(
        "Dense output stats - Mean: {:.6}, Std: {:.6}",
        compute_mean(&dense1_output),
        compute_std(&dense1_output, compute_mean(&dense1_output))
    );

    println!(
        "After BatchNorm - Mean: {:.6}, Std: {:.6}",
        compute_mean(&bn1_output),
        compute_std(&bn1_output, compute_mean(&bn1_output))
    );

    Ok(())
}

pub fn num_features(&self) -> usize

Get the number of features

pub fn momentum(&self) -> f64

Get the momentum value

pub fn eps(&self) -> f64

Get the epsilon value

pub fn is_training(&self) -> bool

Get the training mode

Trait Implementations

impl<F: Clone + Float + Debug + Send + Sync> Clone for BatchNorm<F>

fn clone(&self) -> BatchNorm<F>

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl<F: Debug + Float + Debug + Send + Sync> Debug for BatchNorm<F>

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

Formats the value using the given formatter.

impl<F: Float + Debug + ScalarOperand + Send + Sync + 'static> Layer<F> for BatchNorm<F>

fn as_any(&self) -> &dyn Any

Get the layer as a dyn Any for downcasting

fn as_any_mut(&mut self) -> &mut dyn Any

Get the layer as a mutable dyn Any for downcasting

fn forward(&self, input: &Array<F, IxDyn>) -> Result<Array<F, IxDyn>>

Forward pass of the layer

fn backward( &self, _input: &Array<F, IxDyn>, grad_output: &Array<F, IxDyn>, ) -> Result<Array<F, IxDyn>>

Backward pass of the layer to compute gradients

fn update(&mut self, learning_rate: F) -> Result<()>

Update the layer parameters with the given gradients

fn params(&self) -> Vec<Array<F, IxDyn>>

Get the parameters of the layer

fn gradients(&self) -> Vec<Array<F, IxDyn>>

Get the gradients of the layer parameters

fn set_gradients(&mut self, _gradients: &[Array<F, IxDyn>]) -> Result<()>

Set the gradients of the layer parameters

fn set_params(&mut self, _params: &[Array<F, IxDyn>]) -> Result<()>

Set the parameters of the layer

fn set_training(&mut self, _training: bool)

Set the layer to training mode (true) or evaluation mode (false)

fn is_training(&self) -> bool

Get the current training mode

fn layer_type(&self) -> &str

Get the type of the layer (e.g., “Dense”, “Conv2D”)

fn parameter_count(&self) -> usize

Get the number of trainable parameters in this layer

fn layer_description(&self) -> String

Get a detailed description of this layer

impl<F: Float + Debug + ScalarOperand + Send + Sync + 'static> ParamLayer<F> for BatchNorm<F>

fn get_parameters(&self) -> Vec<&Array<F, IxDyn>>

Get the parameters of the layer as a vector of arrays

fn get_gradients(&self) -> Vec<&Array<F, IxDyn>>

Get the gradients of the parameters

fn set_parameters(&mut self, params: Vec<Array<F, IxDyn>>) -> Result<()>

Set the parameters of the layer