torsh_nn/layers/normalization/
mod.rs

1//! Normalization layers for neural networks
2//!
3//! This module provides a comprehensive collection of normalization techniques used in
4//! deep learning, organized into logical families for better maintainability and clarity.
5//!
6//! ## Module Structure
7//!
8//! The normalization module is organized into specialized sub-modules:
9//!
10//! - **`common`** - Shared utilities, configurations, and helper functions
11//! - **`batch`** - Batch normalization variants (1D, 2D, 3D, synchronized, virtual, renormalization)
12//! - **`instance`** - Instance normalization for all dimensions (1D, 2D, 3D)
13//! - **`layer_group`** - Layer and group normalization techniques
14//! - **`weight_based`** - Weight-based normalization (spectral, weight norm, weight standardization)
15//! - **`advanced`** - Advanced adaptive normalization techniques (switchable norm)
16//!
17//! ## Usage Examples
18//!
19//! ```rust
20//! # use torsh_nn::layers::normalization::{
21//! #     BatchNorm2d, LayerNorm, GroupNorm, InstanceNorm2d, SwitchableNorm2d
22//! # };
23//! # use torsh_core::error::Result;
24//! # fn main() -> Result<()> {
25//! // Create different normalization layers
26//! let batch_norm = BatchNorm2d::new(64)?;
27//! let layer_norm = LayerNorm::new(vec![128])?;
28//! let group_norm = GroupNorm::new(8, 64)?;
29//! let instance_norm = InstanceNorm2d::new(64)?;
30//! let switchable_norm = SwitchableNorm2d::new(64)?;
31//! # Ok(())
32//! # }
33//! ```
34//!
35//! ## Configuration and Customization
36//!
37//! Most normalization layers support custom configuration:
38//!
39//! ```rust
40//! # use torsh_nn::layers::normalization::{BatchNorm2d, NormalizationConfig};
41//! # use torsh_core::error::Result;
42//! # fn main() -> Result<()> {
43//! // Create with custom configuration
44//! let config = NormalizationConfig {
45//!     eps: 1e-6,
46//!     momentum: 0.05,
47//!     ..NormalizationConfig::default()
48//! };
49//!
50//! let batch_norm = BatchNorm2d::with_config(64, config)?;
51//! # Ok(())
52//! # }
53//! ```
54
55// Sub-modules containing different normalization families
56pub mod advanced;
57pub mod batch;
58pub mod common;
59pub mod instance;
60pub mod layer_group;
61pub mod weight_based;
62
63// Re-export common utilities and configurations for convenience
64pub use common::{unbiased_variance, utils, NormalizationConfig, NormalizationStats, RunningStats};
65
66// Re-export all batch normalization variants
67pub use batch::{
68    BatchNorm1d, BatchNorm2d, BatchNorm3d, BatchRenorm2d, BatchRenormSchedule, SyncBatchNorm2d,
69    VirtualBatchNorm2d,
70};
71
72// Re-export instance normalization variants
73pub use instance::{InstanceNorm1d, InstanceNorm2d, InstanceNorm3d};
74
75// Re-export layer and group normalization
76pub use layer_group::{GroupNorm, LayerNorm, RMSNorm};
77
78// Re-export weight-based normalization techniques
79pub use weight_based::{SpectralNorm, WeightNorm, WeightStandardization};
80
81// Re-export advanced normalization techniques
82pub use advanced::SwitchableNorm2d;
83
84// Additional backward compatibility aliases for the most commonly used types
85pub use BatchNorm2d as BatchNorm;
86pub use GroupNorm as GN;
87pub use InstanceNorm2d as InstanceNorm;
88pub use LayerNorm as LN;
89
90/// Normalization layer factory for common configurations
91pub struct NormalizationFactory;
92
93impl NormalizationFactory {
94    /// Create a batch normalization layer for 2D inputs (most common)
95    pub fn batch_norm(num_features: usize) -> torsh_core::error::Result<BatchNorm2d> {
96        BatchNorm2d::new(num_features)
97    }
98
99    /// Create a layer normalization for transformer-like architectures
100    pub fn layer_norm(normalized_shape: Vec<usize>) -> torsh_core::error::Result<LayerNorm> {
101        LayerNorm::new(normalized_shape)
102    }
103
104    /// Create a group normalization layer
105    pub fn group_norm(
106        num_groups: usize,
107        num_channels: usize,
108    ) -> torsh_core::error::Result<GroupNorm> {
109        GroupNorm::new(num_groups, num_channels)
110    }
111
112    /// Create an instance normalization layer for 2D inputs
113    pub fn instance_norm(num_features: usize) -> torsh_core::error::Result<InstanceNorm2d> {
114        InstanceNorm2d::new(num_features)
115    }
116
117    /// Create a switchable normalization layer that adapts between different norms
118    pub fn switchable_norm(num_features: usize) -> torsh_core::error::Result<SwitchableNorm2d> {
119        SwitchableNorm2d::new(num_features)
120    }
121
122    /// Create RMS normalization for transformer models
123    pub fn rms_norm(normalized_shape: Vec<usize>) -> torsh_core::error::Result<RMSNorm> {
124        RMSNorm::new(normalized_shape)
125    }
126
127    /// Create batch normalization optimized for training
128    pub fn batch_norm_training(num_features: usize) -> torsh_core::error::Result<BatchNorm2d> {
129        BatchNorm2d::with_config(num_features, NormalizationConfig::training())
130    }
131
132    /// Create batch normalization optimized for inference
133    pub fn batch_norm_inference(num_features: usize) -> torsh_core::error::Result<BatchNorm2d> {
134        BatchNorm2d::with_config(num_features, NormalizationConfig::inference())
135    }
136
137    /// Create layer normalization without learnable parameters
138    pub fn layer_norm_non_affine(
139        normalized_shape: Vec<usize>,
140    ) -> torsh_core::error::Result<LayerNorm> {
141        LayerNorm::with_config(normalized_shape, NormalizationConfig::non_affine())
142    }
143}
144
145/// Common normalization presets for different architectures
146pub struct NormalizationPresets;
147
148impl NormalizationPresets {
149    /// ResNet-style batch normalization
150    pub fn resnet_batch_norm(num_features: usize) -> torsh_core::error::Result<BatchNorm2d> {
151        BatchNorm2d::with_config(num_features, NormalizationConfig::with_momentum(0.1))
152    }
153
154    /// Transformer-style layer normalization
155    pub fn transformer_layer_norm(hidden_size: usize) -> torsh_core::error::Result<LayerNorm> {
156        LayerNorm::with_config(vec![hidden_size], NormalizationConfig::with_eps(1e-12))
157    }
158
159    /// Style transfer instance normalization (non-affine)
160    pub fn style_transfer_instance_norm(
161        num_features: usize,
162    ) -> torsh_core::error::Result<InstanceNorm2d> {
163        InstanceNorm2d::with_config(num_features, NormalizationConfig::non_affine())
164    }
165
166    /// Group normalization for small batch training
167    pub fn small_batch_group_norm(num_channels: usize) -> torsh_core::error::Result<GroupNorm> {
168        let num_groups = if num_channels >= 32 { 32 } else { num_channels };
169        GroupNorm::new(num_groups, num_channels)
170    }
171
172    /// RMS normalization for LLaMA-style transformers
173    pub fn llama_rms_norm(hidden_size: usize) -> torsh_core::error::Result<RMSNorm> {
174        RMSNorm::with_config(vec![hidden_size], 1e-6, true)
175    }
176
177    /// RMS normalization for GPT-style models
178    pub fn gpt_rms_norm(hidden_size: usize) -> torsh_core::error::Result<RMSNorm> {
179        RMSNorm::with_config(vec![hidden_size], 1e-5, true)
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186    use crate::Module;
187    use torsh_tensor::creation::zeros;
188
189    #[test]
190    fn test_normalization_factory() {
191        // Test factory methods
192        let bn =
193            NormalizationFactory::batch_norm(64).expect("Normalization Factory should succeed");
194        assert_eq!(bn.num_features(), 64);
195
196        let ln = NormalizationFactory::layer_norm(vec![128])
197            .expect("Normalization Factory should succeed");
198        assert_eq!(ln.normalized_shape(), &[128]);
199
200        let gn =
201            NormalizationFactory::group_norm(8, 64).expect("Normalization Factory should succeed");
202        assert_eq!(gn.num_groups(), 8);
203        assert_eq!(gn.num_channels(), 64);
204
205        let inn =
206            NormalizationFactory::instance_norm(32).expect("Normalization Factory should succeed");
207        assert_eq!(inn.num_features(), 32);
208
209        let sn = NormalizationFactory::switchable_norm(16)
210            .expect("Normalization Factory should succeed");
211        assert_eq!(sn.num_features(), 16);
212    }
213
214    #[test]
215    fn test_normalization_presets() {
216        // Test preset configurations
217        let resnet_bn = NormalizationPresets::resnet_batch_norm(64)
218            .expect("Normalization Presets should succeed");
219        assert_eq!(resnet_bn.momentum(), 0.1);
220
221        let transformer_ln = NormalizationPresets::transformer_layer_norm(768)
222            .expect("Normalization Presets should succeed");
223        assert_eq!(transformer_ln.eps(), 1e-12);
224
225        let style_in = NormalizationPresets::style_transfer_instance_norm(64)
226            .expect("Normalization Presets should succeed");
227        // Non-affine should not have weight/bias parameters
228        assert!(style_in.parameters().is_empty());
229
230        let small_batch_gn = NormalizationPresets::small_batch_group_norm(64)
231            .expect("Normalization Presets should succeed");
232        assert_eq!(small_batch_gn.num_groups(), 32);
233    }
234
235    #[test]
236    fn test_module_integration() {
237        // Test that different normalization layers work with sample inputs
238        let input_2d = zeros(&[4, 64]).expect("zeros should succeed");
239        let input_4d = zeros(&[4, 64, 32, 32]).expect("zeros should succeed");
240
241        // Test BatchNorm2d
242        let bn2d = BatchNorm2d::new(64).expect("Batch Norm2d should succeed");
243        assert!(bn2d.forward(&input_4d).is_ok());
244
245        // Test BatchNorm1d
246        let bn1d = BatchNorm1d::new(64).expect("Batch Norm1d should succeed");
247        assert!(bn1d.forward(&input_2d).is_ok());
248
249        // Test LayerNorm
250        let ln = LayerNorm::new(vec![64]).expect("Layer Norm should succeed");
251        assert!(ln.forward(&input_2d).is_ok());
252
253        // Test GroupNorm
254        let gn = GroupNorm::new(8, 64).expect("Group Norm should succeed");
255        assert!(gn.forward(&input_4d).is_ok());
256
257        // Test InstanceNorm2d
258        let in2d = InstanceNorm2d::new(64).expect("Instance Norm2d should succeed");
259        assert!(in2d.forward(&input_4d).is_ok());
260    }
261
262    #[test]
263    fn test_backward_compatibility_aliases() {
264        // Test that aliases work correctly
265        let bn = BatchNorm::new(32).expect("Batch Norm should succeed");
266        assert_eq!(bn.num_features(), 32);
267
268        let ln = LN::new(vec![128]).expect("LN should succeed");
269        assert_eq!(ln.normalized_shape(), &[128]);
270
271        let gn = GN::new(4, 32).expect("GN should succeed");
272        assert_eq!(gn.num_groups(), 4);
273
274        let inn = InstanceNorm::new(16).expect("Instance Norm should succeed");
275        assert_eq!(inn.num_features(), 16);
276    }
277
278    #[test]
279    fn test_configuration_variants() {
280        // Test different configuration variants
281        let training_config = NormalizationConfig::training();
282        assert!(training_config.track_running_stats);
283        assert!(training_config.affine);
284
285        let inference_config = NormalizationConfig::inference();
286        assert!(!inference_config.track_running_stats);
287
288        let non_affine_config = NormalizationConfig::non_affine();
289        assert!(!non_affine_config.affine);
290
291        let custom_eps_config = NormalizationConfig::with_eps(1e-8);
292        assert_eq!(custom_eps_config.eps, 1e-8);
293
294        let custom_momentum_config = NormalizationConfig::with_momentum(0.05);
295        assert_eq!(custom_momentum_config.momentum, 0.05);
296    }
297}
torsh_nn/layers/normalization/mod.rs

torsh_nn/layers/normalization/
mod.rs
