pub trait Module: Send + Sync {
// Required method
fn forward(&self, input: &Variable) -> Variable;
// Provided methods
fn parameters(&self) -> Vec<Parameter> { ... }
fn named_parameters(&self) -> HashMap<String, Parameter> { ... }
fn num_parameters(&self) -> usize { ... }
fn train(&mut self) { ... }
fn eval(&mut self) { ... }
fn set_training(&mut self, _training: bool) { ... }
fn is_training(&self) -> bool { ... }
fn zero_grad(&self) { ... }
fn name(&self) -> &'static str { ... }
}Expand description
Core trait for all neural network modules.
Every layer in Axonml implements this trait, which provides:
- Forward pass computation
- Parameter management
- Training/evaluation mode switching
- Module naming
Required Methods§
Provided Methods§
Sourcefn parameters(&self) -> Vec<Parameter>
fn parameters(&self) -> Vec<Parameter>
Returns all parameters of this module.
This includes parameters from all child modules.
Sourcefn named_parameters(&self) -> HashMap<String, Parameter>
fn named_parameters(&self) -> HashMap<String, Parameter>
Returns named parameters of this module.
Sourcefn num_parameters(&self) -> usize
fn num_parameters(&self) -> usize
Returns the number of trainable parameters.
Sourcefn set_training(&mut self, _training: bool)
fn set_training(&mut self, _training: bool)
Sets the training mode.
Sourcefn is_training(&self) -> bool
fn is_training(&self) -> bool
Returns whether the module is in training mode.