ghostflow_nn/

norm.rs

//! Normalization layers

use ghostflow_core::Tensor;
use crate::module::Module;

/// Batch Normalization for 1D inputs (N, C) or (N, C, L)
pub struct BatchNorm1d {
    #[allow(dead_code)]
    num_features: usize,
    gamma: Tensor,  // scale
    beta: Tensor,   // shift
    running_mean: Tensor,
    running_var: Tensor,
    eps: f32,
    #[allow(dead_code)]
    momentum: f32,
    training: bool,
}

impl BatchNorm1d {
    pub fn new(num_features: usize) -> Self {
        Self::with_params(num_features, 1e-5, 0.1)
    }

    pub fn with_params(num_features: usize, eps: f32, momentum: f32) -> Self {
        BatchNorm1d {
            num_features,
            gamma: Tensor::ones(&[num_features]),
            beta: Tensor::zeros(&[num_features]),
            running_mean: Tensor::zeros(&[num_features]),
            running_var: Tensor::ones(&[num_features]),
            eps,
            momentum,
            training: true,
        }
    }
}

impl Module for BatchNorm1d {
    fn forward(&self, input: &Tensor) -> Tensor {
        let dims = input.dims();
        let data = input.data_f32();
        let gamma = self.gamma.data_f32();
        let beta = self.beta.data_f32();
        
        let batch_size = dims[0];
        let channels = dims[1];
        let spatial_size = if dims.len() > 2 { dims[2] } else { 1 };
        
        let (mean, var) = if self.training {
            // Compute batch statistics
            let mut mean = vec![0.0f32; channels];
            let mut var = vec![0.0f32; channels];
            let n = (batch_size * spatial_size) as f32;
            
            for c in 0..channels {
                let mut sum = 0.0f32;
                for b in 0..batch_size {
                    for s in 0..spatial_size {
                        let idx = b * channels * spatial_size + c * spatial_size + s;
                        sum += data[idx];
                    }
                }
                mean[c] = sum / n;
                
                let mut var_sum = 0.0f32;
                for b in 0..batch_size {
                    for s in 0..spatial_size {
                        let idx = b * channels * spatial_size + c * spatial_size + s;
                        var_sum += (data[idx] - mean[c]).powi(2);
                    }
                }
                var[c] = var_sum / n;
            }
            
            (mean, var)
        } else {
            (self.running_mean.data_f32(), self.running_var.data_f32())
        };
        
        // Normalize
        let mut output = vec![0.0f32; data.len()];
        
        for b in 0..batch_size {
            for c in 0..channels {
                let std = (var[c] + self.eps).sqrt();
                for s in 0..spatial_size {
                    let idx = b * channels * spatial_size + c * spatial_size + s;
                    output[idx] = gamma[c] * (data[idx] - mean[c]) / std + beta[c];
                }
            }
        }
        
        Tensor::from_slice(&output, dims).unwrap()
    }

    fn parameters(&self) -> Vec<Tensor> {
        vec![self.gamma.clone(), self.beta.clone()]
    }

    fn train(&mut self) { self.training = true; }
    fn eval(&mut self) { self.training = false; }
    fn is_training(&self) -> bool { self.training }
}

/// Batch Normalization for 2D inputs (N, C, H, W)
pub struct BatchNorm2d {
    #[allow(dead_code)]
    num_features: usize,
    gamma: Tensor,
    beta: Tensor,
    running_mean: Tensor,
    running_var: Tensor,
    eps: f32,
    #[allow(dead_code)]
    momentum: f32,
    training: bool,
}

impl BatchNorm2d {
    pub fn new(num_features: usize) -> Self {
        Self::with_params(num_features, 1e-5, 0.1)
    }

    pub fn with_params(num_features: usize, eps: f32, momentum: f32) -> Self {
        BatchNorm2d {
            num_features,
            gamma: Tensor::ones(&[num_features]),
            beta: Tensor::zeros(&[num_features]),
            running_mean: Tensor::zeros(&[num_features]),
            running_var: Tensor::ones(&[num_features]),
            eps,
            momentum,
            training: true,
        }
    }
}

impl Module for BatchNorm2d {
    fn forward(&self, input: &Tensor) -> Tensor {
        let dims = input.dims();
        let data = input.data_f32();
        let gamma = self.gamma.data_f32();
        let beta = self.beta.data_f32();
        
        let batch_size = dims[0];
        let channels = dims[1];
        let height = dims[2];
        let width = dims[3];
        let spatial_size = height * width;
        
        let (mean, var) = if self.training {
            let mut mean = vec![0.0f32; channels];
            let mut var = vec![0.0f32; channels];
            let n = (batch_size * spatial_size) as f32;
            
            for c in 0..channels {
                let mut sum = 0.0f32;
                for b in 0..batch_size {
                    for h in 0..height {
                        for w in 0..width {
                            let idx = b * channels * spatial_size + c * spatial_size + h * width + w;
                            sum += data[idx];
                        }
                    }
                }
                mean[c] = sum / n;
                
                let mut var_sum = 0.0f32;
                for b in 0..batch_size {
                    for h in 0..height {
                        for w in 0..width {
                            let idx = b * channels * spatial_size + c * spatial_size + h * width + w;
                            var_sum += (data[idx] - mean[c]).powi(2);
                        }
                    }
                }
                var[c] = var_sum / n;
            }
            
            (mean, var)
        } else {
            (self.running_mean.data_f32(), self.running_var.data_f32())
        };
        
        let mut output = vec![0.0f32; data.len()];
        
        for b in 0..batch_size {
            for c in 0..channels {
                let std = (var[c] + self.eps).sqrt();
                for h in 0..height {
                    for w in 0..width {
                        let idx = b * channels * spatial_size + c * spatial_size + h * width + w;
                        output[idx] = gamma[c] * (data[idx] - mean[c]) / std + beta[c];
                    }
                }
            }
        }
        
        Tensor::from_slice(&output, dims).unwrap()
    }

    fn parameters(&self) -> Vec<Tensor> {
        vec![self.gamma.clone(), self.beta.clone()]
    }

    fn train(&mut self) { self.training = true; }
    fn eval(&mut self) { self.training = false; }
    fn is_training(&self) -> bool { self.training }
}

/// Layer Normalization
pub struct LayerNorm {
    normalized_shape: Vec<usize>,
    gamma: Tensor,
    beta: Tensor,
    eps: f32,
    training: bool,
}

impl LayerNorm {
    pub fn new(normalized_shape: &[usize]) -> Self {
        Self::with_eps(normalized_shape, 1e-5)
    }

    pub fn with_eps(normalized_shape: &[usize], eps: f32) -> Self {
        let numel: usize = normalized_shape.iter().product();
        
        LayerNorm {
            normalized_shape: normalized_shape.to_vec(),
            gamma: Tensor::ones(&[numel]),
            beta: Tensor::zeros(&[numel]),
            eps,
            training: true,
        }
    }
}

impl Module for LayerNorm {
    fn forward(&self, input: &Tensor) -> Tensor {
        let dims = input.dims();
        let data = input.data_f32();
        let gamma = self.gamma.data_f32();
        let beta = self.beta.data_f32();
        
        let norm_size: usize = self.normalized_shape.iter().product();
        let batch_size = data.len() / norm_size;
        
        let mut output = vec![0.0f32; data.len()];
        
        for b in 0..batch_size {
            let start = b * norm_size;
            let end = start + norm_size;
            let slice = &data[start..end];
            
            // Compute mean
            let mean: f32 = slice.iter().sum::<f32>() / norm_size as f32;
            
            // Compute variance
            let var: f32 = slice.iter().map(|&x| (x - mean).powi(2)).sum::<f32>() / norm_size as f32;
            let std = (var + self.eps).sqrt();
            
            // Normalize
            for i in 0..norm_size {
                output[start + i] = gamma[i] * (slice[i] - mean) / std + beta[i];
            }
        }
        
        Tensor::from_slice(&output, dims).unwrap()
    }

    fn parameters(&self) -> Vec<Tensor> {
        vec![self.gamma.clone(), self.beta.clone()]
    }

    fn train(&mut self) { self.training = true; }
    fn eval(&mut self) { self.training = false; }
    fn is_training(&self) -> bool { self.training }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_batchnorm2d() {
        let bn = BatchNorm2d::new(16);
        let input = Tensor::randn(&[2, 16, 8, 8]);
        let output = bn.forward(&input);
        
        assert_eq!(output.dims(), input.dims());
    }

    #[test]
    fn test_layernorm() {
        let ln = LayerNorm::new(&[64]);
        let input = Tensor::randn(&[2, 10, 64]);
        let output = ln.forward(&input);
        
        assert_eq!(output.dims(), input.dims());
    }
}