optirs_core/regularizers/
spatial_dropout.rs

1// Spatial and Feature Dropout regularization
2//
3// This module provides specialized dropout variants that preserve spatial or feature connectivity:
4// - Spatial Dropout: drops entire feature maps (useful for CNNs)
5// - Feature Dropout: drops specific features across all spatial locations
6
7use scirs2_core::ndarray::{Array, Axis, Dimension, ScalarOperand};
8use scirs2_core::numeric::Float;
9use scirs2_core::random::thread_rng;
10use std::fmt::Debug;
11
12use crate::error::{OptimError, Result};
13use crate::regularizers::Regularizer;
14
15/// Spatial Dropout regularizer
16///
17/// Drops entire feature maps instead of individual units. This helps
18/// preserve spatial structure in convolutional neural networks.
19///
20/// # Example
21///
22/// ```
23/// use scirs2_core::ndarray::Array4;
24/// use optirs_core::regularizers::SpatialDropout;
25///
26/// let spatial_dropout = SpatialDropout::new(0.3).expect("SpatialDropout::new succeeds"); // 30% dropout rate
27///
28/// // 4D tensor (batch, channels, height, width)
29/// let features = Array4::<f64>::ones((2, 3, 4, 4));
30///
31/// // During training - drops entire channels
32/// let masked_features = spatial_dropout.apply(&features, true);
33/// ```
34#[derive(Debug, Clone)]
35pub struct SpatialDropout<A: Float> {
36    /// Probability of dropping a channel/feature map
37    dropprob: A,
38    /// Dimension along which to drop (default is 1 for channels)
39    feature_dim: Axis,
40}
41
42impl<A: Float + Debug + ScalarOperand + Send + Sync> SpatialDropout<A> {
43    /// Create a new SpatialDropout regularizer
44    ///
45    /// # Arguments
46    ///
47    /// * `dropprob` - Probability of dropping each feature map (0.0 to 1.0)
48    pub fn new(dropprob: A) -> Result<Self> {
49        if dropprob < A::zero() || dropprob > A::one() {
50            return Err(OptimError::InvalidConfig(
51                "Drop probability must be between 0.0 and 1.0".to_string(),
52            ));
53        }
54
55        Ok(Self {
56            dropprob,
57            feature_dim: Axis(1), // Default to channel dimension
58        })
59    }
60
61    /// Set the dimension along which to drop features
62    pub fn with_feature_dim(mut self, dim: usize) -> Self {
63        self.feature_dim = Axis(dim);
64        self
65    }
66
67    /// Apply spatial dropout to a tensor
68    pub fn apply<D>(&self, features: &Array<A, D>, training: bool) -> Array<A, D>
69    where
70        D: Dimension + scirs2_core::ndarray::RemoveAxis,
71    {
72        if !training || self.dropprob == A::zero() {
73            return features.clone();
74        }
75
76        let keep_prob = A::one() - self.dropprob;
77
78        // Get the size of the feature dimension
79        let feature_size = features.shape()[self.feature_dim.0];
80
81        // Create a mask for each feature map
82        let keep_prob_f64 = keep_prob
83            .to_f64()
84            .expect("SpatialDropout: keep_prob in [0, 1] (validated at construction) fits in f64");
85        let mut rng = thread_rng();
86        let feature_mask: Vec<bool> = (0..feature_size)
87            .map(|_| rng.random_bool(keep_prob_f64))
88            .collect();
89
90        // Apply mask to each feature map
91        let mut result = features.clone();
92        for (idx, &keep) in feature_mask.iter().enumerate() {
93            if !keep {
94                // Drop the entire feature map
95                let mut axis_slice = result.index_axis_mut(self.feature_dim, idx);
96                axis_slice.fill(A::zero());
97            } else {
98                // Scale kept features
99                let mut axis_slice = result.index_axis_mut(self.feature_dim, idx);
100                axis_slice.mapv_inplace(|x| x / keep_prob);
101            }
102        }
103
104        result
105    }
106}
107
108/// Feature Dropout regularizer
109///
110/// Drops specific features across all spatial locations. This is useful when
111/// you want to maintain spatial consistency while dropping features.
112///
113/// # Example
114///
115/// ```
116/// use scirs2_core::ndarray::Array3;
117/// use optirs_core::regularizers::FeatureDropout;
118///
119/// let feature_dropout = FeatureDropout::new(0.5).expect("FeatureDropout::new succeeds"); // 50% dropout rate
120///
121/// // 3D tensor (batch, features, sequence_length)
122/// let features = Array3::<f64>::ones((2, 10, 20));
123///
124/// // During training - drops specific features across all positions
125/// let masked_features = feature_dropout.apply(&features, true);
126/// ```
127#[derive(Debug, Clone)]
128pub struct FeatureDropout<A: Float> {
129    /// Probability of dropping each feature
130    dropprob: A,
131    /// Dimension along which features are located (default is 1)
132    feature_dim: Axis,
133}
134
135impl<A: Float + Debug + ScalarOperand + Send + Sync> FeatureDropout<A> {
136    /// Create a new FeatureDropout regularizer
137    ///
138    /// # Arguments
139    ///
140    /// * `dropprob` - Probability of dropping each feature (0.0 to 1.0)
141    pub fn new(dropprob: A) -> Result<Self> {
142        if dropprob < A::zero() || dropprob > A::one() {
143            return Err(OptimError::InvalidConfig(
144                "Drop probability must be between 0.0 and 1.0".to_string(),
145            ));
146        }
147
148        Ok(Self {
149            dropprob,
150            feature_dim: Axis(1), // Default to feature dimension
151        })
152    }
153
154    /// Set the dimension along which features are located
155    pub fn with_feature_dim(mut self, dim: usize) -> Self {
156        self.feature_dim = Axis(dim);
157        self
158    }
159
160    /// Apply feature dropout to a tensor
161    pub fn apply<D>(&self, features: &Array<A, D>, training: bool) -> Array<A, D>
162    where
163        D: Dimension + scirs2_core::ndarray::RemoveAxis,
164    {
165        if !training || self.dropprob == A::zero() {
166            return features.clone();
167        }
168
169        let keep_prob = A::one() - self.dropprob;
170
171        // Get the size of the feature dimension
172        let feature_size = features.shape()[self.feature_dim.0];
173
174        // Create a consistent mask for each feature
175        let keep_prob_f64 = keep_prob
176            .to_f64()
177            .expect("FeatureDropout: keep_prob in [0, 1] (validated at construction) fits in f64");
178        let mut rng = thread_rng();
179        let feature_mask: Vec<bool> = (0..feature_size)
180            .map(|_| rng.random_bool(keep_prob_f64))
181            .collect();
182
183        // Apply the same mask across all spatial/temporal locations
184        let mut result = features.clone();
185        for (idx, &keep) in feature_mask.iter().enumerate() {
186            if !keep {
187                // Drop this feature everywhere
188                let mut axis_slice = result.index_axis_mut(self.feature_dim, idx);
189                axis_slice.fill(A::zero());
190            } else {
191                // Scale kept features
192                let mut axis_slice = result.index_axis_mut(self.feature_dim, idx);
193                axis_slice.mapv_inplace(|x| x / keep_prob);
194            }
195        }
196
197        result
198    }
199}
200
201// Implement Regularizer trait for SpatialDropout - only for dimensions that support RemoveAxis
202impl<
203        A: Float + Debug + ScalarOperand + Send + Sync,
204        D: Dimension + scirs2_core::ndarray::RemoveAxis + Send + Sync,
205    > Regularizer<A, D> for SpatialDropout<A>
206{
207    fn apply(&self, _params: &Array<A, D>, gradients: &mut Array<A, D>) -> Result<A> {
208        // Apply spatial dropout to gradients during training
209        let masked_gradients = SpatialDropout::apply(self, gradients, true);
210        *gradients = masked_gradients;
211        Ok(A::zero())
212    }
213
214    fn penalty(&self, _params: &Array<A, D>) -> Result<A> {
215        // Spatial dropout doesn't add a penalty term
216        Ok(A::zero())
217    }
218}
219
220// Implement Regularizer trait for FeatureDropout - only for dimensions that support RemoveAxis
221impl<
222        A: Float + Debug + ScalarOperand + Send + Sync,
223        D: Dimension + scirs2_core::ndarray::RemoveAxis + Send + Sync,
224    > Regularizer<A, D> for FeatureDropout<A>
225{
226    fn apply(&self, _params: &Array<A, D>, gradients: &mut Array<A, D>) -> Result<A> {
227        // Apply feature dropout to gradients during training
228        let masked_gradients = FeatureDropout::apply(self, gradients, true);
229        *gradients = masked_gradients;
230        Ok(A::zero())
231    }
232
233    fn penalty(&self, _params: &Array<A, D>) -> Result<A> {
234        // Feature dropout doesn't add a penalty term
235        Ok(A::zero())
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242    use approx::assert_relative_eq;
243    use scirs2_core::ndarray::{array, Ix3};
244
245    #[test]
246    fn test_spatial_dropout_creation() {
247        // Valid creation
248        let sd = SpatialDropout::<f64>::new(0.3)
249            .expect("SpatialDropout::<f64>::new succeeds in test_spatial_dropout_creation");
250        assert_eq!(sd.dropprob, 0.3);
251
252        // Invalid probabilities
253        assert!(SpatialDropout::<f64>::new(-0.1).is_err());
254        assert!(SpatialDropout::<f64>::new(1.1).is_err());
255    }
256
257    #[test]
258    fn test_spatial_dropout_4d() {
259        let sd = SpatialDropout::new(0.5)
260            .expect("SpatialDropout::new succeeds in test_spatial_dropout_4d");
261
262        // Create a 4D tensor (batch, channels, height, width)
263        // Use values that are always non-zero to better test dropout
264        let features = Array::from_shape_fn((2, 4, 3, 3), |(b, c, h, w)| {
265            1.0 + b as f64 + c as f64 * 10.0 + h as f64 * 0.1 + w as f64 * 0.01
266        });
267
268        // Apply spatial dropout
269        let masked = sd.apply(&features, true);
270
271        // Check that entire channels are either kept or dropped
272        for b in 0..2 {
273            for c in 0..4 {
274                let masked_batch = masked.index_axis(Axis(0), b);
275                let channel = masked_batch.index_axis(Axis(0), c);
276                let channel_clone = channel.to_owned();
277                let is_dropped = channel_clone.iter().all(|&x| x.abs() < 1e-10);
278                let is_kept = channel_clone.iter().all(|&x| x.abs() > 1e-10);
279
280                // For dropped channels, all values should be 0
281                // For kept channels, all values should be scaled by 1/keep_prob
282                if is_dropped {
283                    for &val in channel_clone.iter() {
284                        assert_eq!(val, 0.0);
285                    }
286                } else if is_kept {
287                    // Check scaling
288                    let original_batch = features.index_axis(Axis(0), b);
289                    let original_channel = original_batch.index_axis(Axis(0), c);
290                    for ((i, j), &val) in channel_clone.indexed_iter() {
291                        assert_relative_eq!(val, original_channel[[i, j]] * 2.0, epsilon = 1e-10);
292                    }
293                } else {
294                    // Mixed values - this shouldn't happen
295                    println!("Channel {c} in batch {b} has mixed values:");
296                    for val in channel_clone.iter() {
297                        println!("  Value: {val}");
298                    }
299                    panic!("Channel should be entirely dropped or kept");
300                }
301            }
302        }
303    }
304
305    #[test]
306    fn test_feature_dropout_creation() {
307        // Valid creation
308        let fd = FeatureDropout::<f64>::new(0.4)
309            .expect("FeatureDropout::<f64>::new succeeds in test_feature_dropout_creation");
310        assert_eq!(fd.dropprob, 0.4);
311
312        // Invalid probabilities
313        assert!(FeatureDropout::<f64>::new(-0.1).is_err());
314        assert!(FeatureDropout::<f64>::new(1.1).is_err());
315    }
316
317    #[test]
318    fn test_feature_dropout_3d() {
319        let fd = FeatureDropout::new(0.5)
320            .expect("FeatureDropout::new succeeds in test_feature_dropout_3d");
321
322        // Create a 3D tensor (batch, features, sequence)
323        let features = Array::from_shape_fn((2, 5, 10), |(_b, f, s)| f as f64 + s as f64);
324
325        // Apply feature dropout
326        let masked = fd.apply(&features, true);
327
328        // Check that features are consistently dropped across all positions
329        for f in 0..5 {
330            let first_batch = masked.index_axis(Axis(0), 0);
331            let first_batch_feature = first_batch.index_axis(Axis(0), f);
332            let first_batch_clone = first_batch_feature.to_owned();
333            let is_dropped = first_batch_clone.iter().all(|&x| x == 0.0);
334
335            // Check consistency across batches and positions
336            for b in 0..2 {
337                let batch = masked.index_axis(Axis(0), b);
338                let feature_slice = batch.index_axis(Axis(0), f);
339                let feature_clone = feature_slice.to_owned();
340                let all_dropped = feature_clone.iter().all(|&x| x == 0.0);
341                assert_eq!(
342                    is_dropped, all_dropped,
343                    "Feature dropout should be consistent"
344                );
345
346                if !all_dropped {
347                    // Check scaling
348                    let original_batch = features.index_axis(Axis(0), b);
349                    let original_slice = original_batch.index_axis(Axis(0), f);
350                    for (i, &val) in feature_clone.iter().enumerate() {
351                        assert_relative_eq!(val, original_slice[i] * 2.0, epsilon = 1e-10);
352                    }
353                }
354            }
355        }
356    }
357
358    #[test]
359    fn test_inference_mode() {
360        let sd =
361            SpatialDropout::new(0.5).expect("SpatialDropout::new succeeds in test_inference_mode");
362        let fd =
363            FeatureDropout::new(0.5).expect("FeatureDropout::new succeeds in test_inference_mode");
364
365        let features = array![[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]];
366
367        // During inference, features should remain unchanged
368        let sd_inference = sd.apply(&features, false);
369        let fd_inference = fd.apply(&features, false);
370
371        assert_eq!(features, sd_inference);
372        assert_eq!(features, fd_inference);
373    }
374
375    #[test]
376    fn test_regularizer_trait() {
377        let sd = SpatialDropout::new(0.3)
378            .expect("SpatialDropout::new succeeds in test_regularizer_trait");
379        let params = array![[[1.0, 2.0], [3.0, 4.0]]];
380        let mut gradient = array![[[0.1, 0.2], [0.3, 0.4]]];
381
382        // Test Regularizer trait
383        let penalty = sd
384            .penalty(&params)
385            .expect("sd.penalty succeeds in test_regularizer_trait");
386        assert_eq!(penalty, 0.0);
387
388        let _penalty_apply = sd.apply(&params, true);
389        let penalty_reg =
390            <SpatialDropout<f64> as Regularizer<f64, Ix3>>::apply(&sd, &params, &mut gradient)
391                .expect("SpatialDropout as Regularizer::apply succeeds in test_regularizer_trait");
392        assert_eq!(penalty_reg, 0.0);
393
394        // Gradient should be modified
395        let is_modified = gradient != array![[[0.1, 0.2], [0.3, 0.4]]];
396        assert!(is_modified || gradient == array![[[0.1, 0.2], [0.3, 0.4]]]);
397    }
398}
optirs_core/regularizers/spatial_dropout.rs

optirs_core/regularizers/
spatial_dropout.rs
