pub trait Module: Send + Sync {
// Required method
fn forward(&self, input: &Tensor) -> Tensor;
// Provided methods
fn parameters(&self) -> Vec<&Tensor> { ... }
fn parameters_mut(&mut self) -> Vec<&mut Tensor> { ... }
fn train(&mut self) { ... }
fn eval(&mut self) { ... }
fn training(&self) -> bool { ... }
fn zero_grad(&mut self) { ... }
fn num_parameters(&self) -> usize { ... }
}Expand description
Base trait for all neural network modules.
Every layer, activation function, and container implements this trait, providing a uniform interface for:
- Forward computation
- Parameter access (for optimizers)
- Training/evaluation mode switching
§Example
use aprender::nn::{Module, Linear};
use aprender::autograd::Tensor;
let layer = Linear::new(10, 5);
let x = Tensor::randn(&[32, 10]);
let output = layer.forward(&x); // [32, 5]
// Access parameters for gradient descent
for param in layer.parameters() {
println!("Shape: {:?}", param.shape());
}Required Methods§
Provided Methods§
Sourcefn parameters(&self) -> Vec<&Tensor>
fn parameters(&self) -> Vec<&Tensor>
Get references to all learnable parameters.
Used by optimizers to iterate over parameters for gradient updates. Parameters are returned in a deterministic order.
Sourcefn parameters_mut(&mut self) -> Vec<&mut Tensor>
fn parameters_mut(&mut self) -> Vec<&mut Tensor>
Get mutable references to all learnable parameters.
Used by optimizers to update parameters in-place.
Sourcefn train(&mut self)
fn train(&mut self)
Set the module to training mode.
This affects layers like Dropout (active during training) and BatchNorm (uses batch statistics during training).
Sourcefn eval(&mut self)
fn eval(&mut self)
Set the module to evaluation mode.
This affects layers like Dropout (disabled during eval) and BatchNorm (uses running statistics during eval).
Sourcefn zero_grad(&mut self)
fn zero_grad(&mut self)
Zero out gradients for all parameters.
Should be called before each training iteration.
Sourcefn num_parameters(&self) -> usize
fn num_parameters(&self) -> usize
Get the number of learnable parameters.