Skip to main content

Module

axonml_distributed::prelude

Trait Module 

Source 
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§

Source

fn forward(&self, input: &Variable) -> Variable

Performs the forward pass.

§Arguments
  • input - Input variable
§Returns

Output variable after applying this module’s transformation.

Provided Methods§

Source

fn parameters(&self) -> Vec<Parameter>

Returns all parameters of this module.

This includes parameters from all child modules.

Source

fn named_parameters(&self) -> HashMap<String, Parameter>

Returns named parameters of this module.

Source

fn num_parameters(&self) -> usize

Returns the number of trainable parameters.

Source

fn train(&mut self)

Sets the module to training mode.

Source

fn eval(&mut self)

Sets the module to evaluation mode.

Source

fn set_training(&mut self, _training: bool)

Sets the training mode.

Source

fn is_training(&self) -> bool

Returns whether the module is in training mode.

Source

fn zero_grad(&self)

Zeros all gradients of parameters.

Source

fn name(&self) -> &'static str

Returns the module name for debugging.

Implementors§