torsh_nn/container/
basic.rs

1//! Basic container modules for organizing layers
2
3use crate::{Module, ModuleBase, Parameter};
4use parking_lot::RwLock;
5use std::sync::Arc;
6use torsh_core::device::DeviceType;
7use torsh_core::error::{Result, TorshError};
8use torsh_tensor::Tensor;
9
10// Conditional imports for std/no_std compatibility
11#[cfg(feature = "std")]
12use std::{boxed::Box, collections::HashMap, vec::Vec};
13
14#[cfg(not(feature = "std"))]
15use alloc::{boxed::Box, vec::Vec};
16
17#[cfg(not(feature = "std"))]
18use hashbrown::HashMap;
19
20/// Sequential container
21pub struct Sequential {
22    base: ModuleBase,
23    modules: Vec<Box<dyn Module>>,
24}
25
26impl std::fmt::Debug for Sequential {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        f.debug_struct("Sequential")
29            .field("modules_count", &self.modules.len())
30            .field("training", &self.base.training())
31            .finish()
32    }
33}
34
35impl Sequential {
36    /// Create a new sequential container
37    pub fn new() -> Self {
38        Self {
39            base: ModuleBase::new(),
40            modules: Vec::new(),
41        }
42    }
43
44    /// Add a module to the sequential container
45    #[allow(clippy::should_implement_trait)]
46    pub fn add<M: Module + 'static>(mut self, module: M) -> Self {
47        self.modules.push(Box::new(module));
48        self
49    }
50
51    /// Add a function as a module
52    pub fn add_fn<F>(mut self, f: F) -> Self
53    where
54        F: Fn(&Tensor) -> Result<Tensor> + Send + Sync + 'static,
55    {
56        self.modules.push(Box::new(FunctionModule::new(f)));
57        self
58    }
59}
60
61impl Default for Sequential {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl Module for Sequential {
68    fn forward(&self, input: &Tensor) -> Result<Tensor> {
69        let mut output = input.clone();
70
71        for module in &self.modules {
72            output = module.forward(&output)?;
73        }
74
75        Ok(output)
76    }
77
78    fn parameters(&self) -> HashMap<String, Parameter> {
79        let mut params = HashMap::new();
80
81        for (i, module) in self.modules.iter().enumerate() {
82            for (name, param) in module.parameters() {
83                params.insert(format!("{}.{}", i, name), param);
84            }
85        }
86
87        params
88    }
89
90    fn named_parameters(&self) -> HashMap<String, Parameter> {
91        let mut params = HashMap::new();
92
93        for (i, module) in self.modules.iter().enumerate() {
94            for (name, param) in module.named_parameters() {
95                params.insert(format!("{}.{}", i, name), param);
96            }
97        }
98
99        params
100    }
101
102    /// Every child's buffers, in child order.
103    ///
104    /// Without this override the container answered the *trait default*
105    /// (`Vec::new()`), so a `Sequential` holding a `BatchNorm` reported zero
106    /// buffers and its `state_dict()` silently dropped every running statistic
107    /// — a saved model reloaded with freshly initialized statistics and
108    /// evaluated differently, with no error anywhere. See
109    /// `tests/hardening_nn_state_dict.rs`.
110    fn buffers(&self) -> Vec<Arc<RwLock<Tensor>>> {
111        self.modules
112            .iter()
113            .flat_map(|module| module.buffers())
114            .collect()
115    }
116
117    /// The same buffers as [`Self::buffers`], keyed `"{index}.{name}"` — the
118    /// key scheme [`Self::named_parameters`] already uses, so the parameter and
119    /// buffer halves of a checkpoint agree on how to address a child.
120    fn named_buffers(&self) -> HashMap<String, Arc<RwLock<Tensor>>> {
121        let mut buffers = HashMap::new();
122
123        for (i, module) in self.modules.iter().enumerate() {
124            for (name, buffer) in module.named_buffers() {
125                buffers.insert(format!("{}.{}", i, name), buffer);
126            }
127        }
128
129        buffers
130    }
131
132    fn train(&mut self) {
133        self.base.set_training(true);
134        for module in &mut self.modules {
135            module.train();
136        }
137    }
138
139    fn eval(&mut self) {
140        self.base.set_training(false);
141        for module in &mut self.modules {
142            module.eval();
143        }
144    }
145
146    fn training(&self) -> bool {
147        self.base.training()
148    }
149
150    fn set_training(&mut self, training: bool) {
151        self.base.set_training(training);
152        for module in &mut self.modules {
153            module.set_training(training);
154        }
155    }
156
157    fn to_device(&mut self, device: DeviceType) -> Result<()> {
158        self.base.to_device(device)?;
159        for module in &mut self.modules {
160            module.to_device(device)?;
161        }
162        Ok(())
163    }
164
165    fn children(&self) -> Vec<&dyn Module> {
166        self.modules.iter().map(|m| m.as_ref()).collect()
167    }
168
169    /// The same children as [`Self::children`], addressed by their index.
170    ///
171    /// Every name-carrying recursion in the trait
172    /// (`all_named_parameters`, `all_named_buffers`, `named_modules`) walks
173    /// `named_children()`, not `children()`. Overriding only the latter left
174    /// the container invisible to all three.
175    fn named_children(&self) -> Vec<(String, &dyn Module)> {
176        self.modules
177            .iter()
178            .enumerate()
179            .map(|(i, module)| (i.to_string(), module.as_ref()))
180            .collect()
181    }
182}
183
184/// ModuleList container
185pub struct ModuleList {
186    base: ModuleBase,
187    modules: Vec<Box<dyn Module>>,
188}
189
190impl std::fmt::Debug for ModuleList {
191    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
192        f.debug_struct("ModuleList")
193            .field("modules_count", &self.modules.len())
194            .field("training", &self.base.training())
195            .finish()
196    }
197}
198
199impl ModuleList {
200    pub fn new() -> Self {
201        Self {
202            base: ModuleBase::new(),
203            modules: Vec::new(),
204        }
205    }
206
207    pub fn len(&self) -> usize {
208        self.modules.len()
209    }
210
211    pub fn is_empty(&self) -> bool {
212        self.modules.is_empty()
213    }
214
215    pub fn push<M: Module + 'static>(&mut self, module: M) {
216        self.modules.push(Box::new(module));
217    }
218
219    pub fn extend<I>(&mut self, modules: I)
220    where
221        I: IntoIterator<Item = Box<dyn Module>>,
222    {
223        self.modules.extend(modules);
224    }
225
226    pub fn get(&self, _index: usize) -> Option<&dyn Module> {
227        self.modules.get(_index).map(|m| m.as_ref())
228    }
229
230    pub fn get_mut(&mut self, index: usize) -> Option<&mut (dyn Module + '_)> {
231        if let Some(m) = self.modules.get_mut(index) {
232            Some(&mut **m)
233        } else {
234            None
235        }
236    }
237}
238
239impl Default for ModuleList {
240    fn default() -> Self {
241        Self::new()
242    }
243}
244
245impl Module for ModuleList {
246    fn forward(&self, _input: &Tensor) -> Result<Tensor> {
247        // ModuleList doesn't define forward pass - each module should be called individually
248        Err(TorshError::InvalidArgument(
249            "ModuleList doesn't define forward pass".to_string(),
250        ))
251    }
252
253    fn parameters(&self) -> HashMap<String, Parameter> {
254        let mut params = HashMap::new();
255
256        for (i, module) in self.modules.iter().enumerate() {
257            for (name, param) in module.parameters() {
258                params.insert(format!("{}.{}", i, name), param);
259            }
260        }
261
262        params
263    }
264
265    fn named_parameters(&self) -> HashMap<String, Parameter> {
266        let mut params = HashMap::new();
267
268        for (i, module) in self.modules.iter().enumerate() {
269            for (name, param) in module.named_parameters() {
270                params.insert(format!("{}.{}", i, name), param);
271            }
272        }
273
274        params
275    }
276
277    /// Every child's buffers, in child order. See [`Sequential::buffers`].
278    fn buffers(&self) -> Vec<Arc<RwLock<Tensor>>> {
279        self.modules
280            .iter()
281            .flat_map(|module| module.buffers())
282            .collect()
283    }
284
285    /// The same buffers as [`Self::buffers`], keyed `"{index}.{name}"`,
286    /// mirroring [`Self::named_parameters`].
287    fn named_buffers(&self) -> HashMap<String, Arc<RwLock<Tensor>>> {
288        let mut buffers = HashMap::new();
289
290        for (i, module) in self.modules.iter().enumerate() {
291            for (name, buffer) in module.named_buffers() {
292                buffers.insert(format!("{}.{}", i, name), buffer);
293            }
294        }
295
296        buffers
297    }
298
299    fn train(&mut self) {
300        self.base.set_training(true);
301        for module in &mut self.modules {
302            module.train();
303        }
304    }
305
306    fn eval(&mut self) {
307        self.base.set_training(false);
308        for module in &mut self.modules {
309            module.eval();
310        }
311    }
312
313    fn training(&self) -> bool {
314        self.base.training()
315    }
316
317    fn set_training(&mut self, training: bool) {
318        self.base.set_training(training);
319        for module in &mut self.modules {
320            module.set_training(training);
321        }
322    }
323
324    fn to_device(&mut self, device: DeviceType) -> Result<()> {
325        self.base.to_device(device)?;
326        for module in &mut self.modules {
327            module.to_device(device)?;
328        }
329        Ok(())
330    }
331
332    fn children(&self) -> Vec<&dyn Module> {
333        self.modules.iter().map(|m| m.as_ref()).collect()
334    }
335
336    /// The same children as [`Self::children`], addressed by their index. See
337    /// [`Sequential::named_children`] for why the trait needs both.
338    fn named_children(&self) -> Vec<(String, &dyn Module)> {
339        self.modules
340            .iter()
341            .enumerate()
342            .map(|(i, module)| (i.to_string(), module.as_ref()))
343            .collect()
344    }
345}
346
347/// ModuleDict container
348pub struct ModuleDict {
349    base: ModuleBase,
350    modules: HashMap<String, Box<dyn Module>>,
351}
352
353impl std::fmt::Debug for ModuleDict {
354    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
355        f.debug_struct("ModuleDict")
356            .field("modules_count", &self.modules.len())
357            .field("training", &self.base.training())
358            .finish()
359    }
360}
361
362impl ModuleDict {
363    pub fn new() -> Self {
364        Self {
365            base: ModuleBase::new(),
366            modules: HashMap::new(),
367        }
368    }
369
370    pub fn len(&self) -> usize {
371        self.modules.len()
372    }
373
374    pub fn is_empty(&self) -> bool {
375        self.modules.is_empty()
376    }
377
378    pub fn insert<M: Module + 'static>(&mut self, key: String, module: M) {
379        self.modules.insert(key, Box::new(module));
380    }
381
382    pub fn get(&self, key: &str) -> Option<&dyn Module> {
383        self.modules.get(key).map(|m| m.as_ref())
384    }
385
386    pub fn get_mut(&mut self, key: &str) -> Option<&mut (dyn Module + '_)> {
387        if let Some(m) = self.modules.get_mut(key) {
388            Some(&mut **m)
389        } else {
390            None
391        }
392    }
393
394    pub fn keys(&self) -> impl Iterator<Item = &String> {
395        self.modules.keys()
396    }
397}
398
399impl Default for ModuleDict {
400    fn default() -> Self {
401        Self::new()
402    }
403}
404
405impl Module for ModuleDict {
406    fn forward(&self, _input: &Tensor) -> Result<Tensor> {
407        // ModuleDict doesn't define forward pass - each module should be called individually
408        Err(TorshError::InvalidArgument(
409            "ModuleDict doesn't define forward pass".to_string(),
410        ))
411    }
412
413    fn parameters(&self) -> HashMap<String, Parameter> {
414        let mut params = HashMap::new();
415
416        for (module_name, module) in &self.modules {
417            for (param_name, param) in module.parameters() {
418                params.insert(format!("{}.{}", module_name, param_name), param);
419            }
420        }
421
422        params
423    }
424
425    fn named_parameters(&self) -> HashMap<String, Parameter> {
426        let mut params = HashMap::new();
427
428        for (module_name, module) in &self.modules {
429            for (param_name, param) in module.named_parameters() {
430                params.insert(format!("{}.{}", module_name, param_name), param);
431            }
432        }
433
434        params
435    }
436
437    /// Every child's buffers. See [`Sequential::buffers`].
438    fn buffers(&self) -> Vec<Arc<RwLock<Tensor>>> {
439        self.modules
440            .values()
441            .flat_map(|module| module.buffers())
442            .collect()
443    }
444
445    /// The same buffers as [`Self::buffers`], keyed `"{key}.{name}"`,
446    /// mirroring [`Self::named_parameters`].
447    fn named_buffers(&self) -> HashMap<String, Arc<RwLock<Tensor>>> {
448        let mut buffers = HashMap::new();
449
450        for (module_name, module) in &self.modules {
451            for (buffer_name, buffer) in module.named_buffers() {
452                buffers.insert(format!("{}.{}", module_name, buffer_name), buffer);
453            }
454        }
455
456        buffers
457    }
458
459    fn train(&mut self) {
460        self.base.set_training(true);
461        for module in self.modules.values_mut() {
462            module.train();
463        }
464    }
465
466    fn eval(&mut self) {
467        self.base.set_training(false);
468        for module in self.modules.values_mut() {
469            module.eval();
470        }
471    }
472
473    fn training(&self) -> bool {
474        self.base.training()
475    }
476
477    fn set_training(&mut self, training: bool) {
478        self.base.set_training(training);
479        for module in self.modules.values_mut() {
480            module.set_training(training);
481        }
482    }
483
484    fn to_device(&mut self, device: DeviceType) -> Result<()> {
485        self.base.to_device(device)?;
486        for module in self.modules.values_mut() {
487            module.to_device(device)?;
488        }
489        Ok(())
490    }
491
492    fn children(&self) -> Vec<&dyn Module> {
493        self.modules.values().map(|m| m.as_ref()).collect()
494    }
495
496    /// The same children as [`Self::children`], addressed by their dictionary
497    /// key. See [`Sequential::named_children`] for why the trait needs both.
498    fn named_children(&self) -> Vec<(String, &dyn Module)> {
499        self.modules
500            .iter()
501            .map(|(name, module)| (name.clone(), module.as_ref()))
502            .collect()
503    }
504}
505
506/// Function module wrapper
507pub struct FunctionModule<F>
508where
509    F: Fn(&Tensor) -> Result<Tensor> + Send + Sync,
510{
511    base: ModuleBase,
512    func: F,
513}
514
515impl<F> FunctionModule<F>
516where
517    F: Fn(&Tensor) -> Result<Tensor> + Send + Sync,
518{
519    pub fn new(func: F) -> Self {
520        Self {
521            base: ModuleBase::new(),
522            func,
523        }
524    }
525}
526
527impl<F> Module for FunctionModule<F>
528where
529    F: Fn(&Tensor) -> Result<Tensor> + Send + Sync,
530{
531    fn forward(&self, input: &Tensor) -> Result<Tensor> {
532        (self.func)(input)
533    }
534
535    fn parameters(&self) -> HashMap<String, Parameter> {
536        HashMap::new()
537    }
538
539    fn named_parameters(&self) -> HashMap<String, Parameter> {
540        HashMap::new()
541    }
542
543    fn train(&mut self) {
544        self.base.set_training(true);
545    }
546
547    fn eval(&mut self) {
548        self.base.set_training(false);
549    }
550
551    fn training(&self) -> bool {
552        self.base.training()
553    }
554
555    fn set_training(&mut self, training: bool) {
556        self.base.set_training(training);
557    }
558
559    fn to_device(&mut self, device: DeviceType) -> Result<()> {
560        self.base.to_device(device)
561    }
562}
torsh_nn/container/basic.rs

torsh_nn/container/
basic.rs
