torsh_nn/layers/normalization/instance/
mod.rs

1//! Instance normalization layers
2//!
3//! Instance normalization normalizes each sample independently across spatial dimensions.
4//! This is particularly useful for style transfer and generative models where batch
5//! statistics may not be meaningful.
6
7use crate::{Module, ModuleBase, Parameter};
8use torsh_core::device::DeviceType;
9use torsh_core::error::Result;
10use torsh_tensor::{creation::*, Tensor};
11
12use super::common::{utils, NormalizationConfig};
13
14// Conditional imports for std/no_std compatibility
15#[cfg(feature = "std")]
16use std::collections::HashMap;
17
18#[cfg(not(feature = "std"))]
19use hashbrown::HashMap;
20
21/// Pull the affine parameters out of a module base, if the layer is affine.
22///
23/// Returns clones of the parameter tensors, which share their storage *and*
24/// their gradient slot with the registered `Parameter`, so `backward()`
25/// accumulates straight into the module's gamma/beta.
26fn affine_tensors(base: &ModuleBase, affine: bool) -> (Option<Tensor>, Option<Tensor>) {
27    if !affine {
28        return (None, None);
29    }
30    let weight = base
31        .parameters
32        .get("weight")
33        .map(|p| p.tensor().read().clone());
34    let bias = base
35        .parameters
36        .get("bias")
37        .map(|p| p.tensor().read().clone());
38    (weight, bias)
39}
40
41/// 1D instance normalization layer, for `(N, C)` inputs.
42///
43/// # Degenerate by construction
44///
45/// This layer's rank contract is `(N, C)`, not PyTorch's `(N, C, L)`: there is
46/// no spatial axis left to normalize over, so each `(sample, channel)` statistic
47/// is taken over a *single* element. The mean is the element itself, the
48/// variance is zero, and the output is therefore exactly `bias`, constant in the
49/// input — `d output / d input` is identically zero.
50///
51/// That is what the arithmetic says, and it is now what `backward()` reports.
52/// While the mean was a detached constant the layer instead claimed a gradient
53/// of `gamma / sqrt(eps)` (about `316 * gamma` at the default epsilon), which
54/// was pure noise. Callers that want a real 1-D instance norm should feed an
55/// `(N, C, L)` tensor to [`crate::functional::instance_norm`].
56pub struct InstanceNorm1d {
57    base: ModuleBase,
58    num_features: usize,
59    config: NormalizationConfig,
60}
61
62impl InstanceNorm1d {
63    pub fn new(num_features: usize) -> Result<Self> {
64        Self::with_config(num_features, NormalizationConfig::default())
65    }
66
67    pub fn with_config(num_features: usize, config: NormalizationConfig) -> Result<Self> {
68        let mut base = ModuleBase::new();
69
70        // Initialize parameters if affine
71        if config.affine {
72            let weight = ones(&[num_features])?;
73            let bias = zeros(&[num_features])?;
74            base.register_parameter("weight".to_string(), Parameter::new(weight));
75            base.register_parameter("bias".to_string(), Parameter::new(bias));
76        }
77
78        Ok(Self {
79            base,
80            num_features,
81            config,
82        })
83    }
84
85    pub fn num_features(&self) -> usize {
86        self.num_features
87    }
88
89    pub fn eps(&self) -> f32 {
90        self.config.eps
91    }
92}
93
94impl Module for InstanceNorm1d {
95    fn forward(&self, input: &Tensor) -> Result<Tensor> {
96        let input_shape = input.shape();
97        let dims = input_shape.dims();
98
99        if dims.len() != 2 {
100            return Err(torsh_core::error::TorshError::InvalidShape(format!(
101                "InstanceNorm1d expects 2D input (N, C), got shape {:?}",
102                dims
103            )));
104        }
105
106        if dims[1] != self.num_features {
107            return Err(torsh_core::error::TorshError::InvalidShape(format!(
108                "Expected {} features, got {}",
109                self.num_features, dims[1]
110            )));
111        }
112
113        let (weight, bias) = affine_tensors(&self.base, self.config.affine);
114        utils::instance_normalize(input, weight.as_ref(), bias.as_ref(), self.config.eps)
115    }
116
117    fn parameters(&self) -> HashMap<String, Parameter> {
118        self.base.named_parameters()
119    }
120
121    fn named_parameters(&self) -> HashMap<String, Parameter> {
122        self.base.named_parameters()
123    }
124
125    fn training(&self) -> bool {
126        self.base.training()
127    }
128
129    fn train(&mut self) {
130        self.base.set_training(true);
131    }
132
133    fn eval(&mut self) {
134        self.base.set_training(false);
135    }
136
137    fn to_device(&mut self, device: DeviceType) -> Result<()> {
138        self.base.to_device(device)
139    }
140}
141
142/// 2D instance normalization layer
143pub struct InstanceNorm2d {
144    base: ModuleBase,
145    num_features: usize,
146    config: NormalizationConfig,
147}
148
149impl InstanceNorm2d {
150    pub fn new(num_features: usize) -> Result<Self> {
151        Self::with_config(num_features, NormalizationConfig::default())
152    }
153
154    pub fn with_config(num_features: usize, config: NormalizationConfig) -> Result<Self> {
155        let mut base = ModuleBase::new();
156
157        // Initialize parameters if affine
158        if config.affine {
159            let weight = ones(&[num_features])?;
160            let bias = zeros(&[num_features])?;
161            base.register_parameter("weight".to_string(), Parameter::new(weight));
162            base.register_parameter("bias".to_string(), Parameter::new(bias));
163        }
164
165        Ok(Self {
166            base,
167            num_features,
168            config,
169        })
170    }
171
172    pub fn num_features(&self) -> usize {
173        self.num_features
174    }
175
176    pub fn eps(&self) -> f32 {
177        self.config.eps
178    }
179}
180
181impl Module for InstanceNorm2d {
182    fn forward(&self, input: &Tensor) -> Result<Tensor> {
183        let input_shape = input.shape();
184        let dims = input_shape.dims();
185
186        if dims.len() != 4 {
187            return Err(torsh_core::error::TorshError::InvalidShape(format!(
188                "InstanceNorm2d expects 4D input (N, C, H, W), got shape {:?}",
189                dims
190            )));
191        }
192
193        if dims[1] != self.num_features {
194            return Err(torsh_core::error::TorshError::InvalidShape(format!(
195                "Expected {} features, got {}",
196                self.num_features, dims[1]
197            )));
198        }
199
200        let (weight, bias) = affine_tensors(&self.base, self.config.affine);
201        utils::instance_normalize(input, weight.as_ref(), bias.as_ref(), self.config.eps)
202    }
203
204    fn parameters(&self) -> HashMap<String, Parameter> {
205        self.base.named_parameters()
206    }
207
208    fn named_parameters(&self) -> HashMap<String, Parameter> {
209        self.base.named_parameters()
210    }
211
212    fn training(&self) -> bool {
213        self.base.training()
214    }
215
216    fn train(&mut self) {
217        self.base.set_training(true);
218    }
219
220    fn eval(&mut self) {
221        self.base.set_training(false);
222    }
223
224    fn to_device(&mut self, device: DeviceType) -> Result<()> {
225        self.base.to_device(device)
226    }
227}
228
229/// 3D instance normalization layer
230pub struct InstanceNorm3d {
231    base: ModuleBase,
232    num_features: usize,
233    config: NormalizationConfig,
234}
235
236impl InstanceNorm3d {
237    pub fn new(num_features: usize) -> Result<Self> {
238        Self::with_config(num_features, NormalizationConfig::default())
239    }
240
241    pub fn with_config(num_features: usize, config: NormalizationConfig) -> Result<Self> {
242        let mut base = ModuleBase::new();
243
244        // Initialize parameters if affine
245        if config.affine {
246            let weight = ones(&[num_features])?;
247            let bias = zeros(&[num_features])?;
248            base.register_parameter("weight".to_string(), Parameter::new(weight));
249            base.register_parameter("bias".to_string(), Parameter::new(bias));
250        }
251
252        Ok(Self {
253            base,
254            num_features,
255            config,
256        })
257    }
258
259    pub fn num_features(&self) -> usize {
260        self.num_features
261    }
262
263    pub fn eps(&self) -> f32 {
264        self.config.eps
265    }
266}
267
268impl Module for InstanceNorm3d {
269    fn forward(&self, input: &Tensor) -> Result<Tensor> {
270        let input_shape = input.shape();
271        let dims = input_shape.dims();
272
273        if dims.len() != 5 {
274            return Err(torsh_core::error::TorshError::InvalidShape(format!(
275                "InstanceNorm3d expects 5D input (N, C, D, H, W), got shape {:?}",
276                dims
277            )));
278        }
279
280        if dims[1] != self.num_features {
281            return Err(torsh_core::error::TorshError::InvalidShape(format!(
282                "Expected {} features, got {}",
283                self.num_features, dims[1]
284            )));
285        }
286
287        let (weight, bias) = affine_tensors(&self.base, self.config.affine);
288        utils::instance_normalize(input, weight.as_ref(), bias.as_ref(), self.config.eps)
289    }
290
291    fn parameters(&self) -> HashMap<String, Parameter> {
292        self.base.named_parameters()
293    }
294
295    fn named_parameters(&self) -> HashMap<String, Parameter> {
296        self.base.named_parameters()
297    }
298
299    fn training(&self) -> bool {
300        self.base.training()
301    }
302
303    fn train(&mut self) {
304        self.base.set_training(true);
305    }
306
307    fn eval(&mut self) {
308        self.base.set_training(false);
309    }
310
311    fn to_device(&mut self, device: DeviceType) -> Result<()> {
312        self.base.to_device(device)
313    }
314}
315
316// Re-export the instance normalization components (already defined in this module)
317
318#[cfg(test)]
319mod tests {
320    use super::*;
321
322    #[test]
323    fn test_instance_norm_2d_creation() {
324        let instance_norm = InstanceNorm2d::new(64).expect("Instance Norm2d should succeed");
325        assert_eq!(instance_norm.num_features(), 64);
326        assert_eq!(instance_norm.eps(), 1e-5);
327    }
328
329    #[test]
330    fn test_instance_norm_2d_shape_validation() {
331        let instance_norm = InstanceNorm2d::new(3).expect("Instance Norm2d should succeed");
332
333        // Valid input
334        let input = zeros(&[2, 3, 32, 32]).expect("zeros should succeed");
335        assert!(instance_norm.forward(&input).is_ok());
336
337        // Invalid dimensions
338        let input_3d = zeros(&[2, 3, 32]).expect("zeros should succeed");
339        assert!(instance_norm.forward(&input_3d).is_err());
340
341        // Wrong number of channels
342        let input_wrong_channels = zeros(&[2, 4, 32, 32]).expect("zeros should succeed");
343        assert!(instance_norm.forward(&input_wrong_channels).is_err());
344    }
345}
torsh_nn/layers/normalization/instance/mod.rs

torsh_nn/layers/normalization/instance/
mod.rs
