Trait burn::module::Module

pub trait Module<B>: Clone + Send + Debugwhere
    B: Backend,{
    type Record: Record<B>;

    // Required methods
    fn collect_devices(
        &self,
        devices: Vec<<B as Backend>::Device>
    ) -> Vec<<B as Backend>::Device>;
    fn fork(self, device: &<B as Backend>::Device) -> Self;
    fn to_device(self, device: &<B as Backend>::Device) -> Self;
    fn visit<Visitor>(&self, visitor: &mut Visitor)
       where Visitor: ModuleVisitor<B>;
    fn map<Mapper>(self, mapper: &mut Mapper) -> Self
       where Mapper: ModuleMapper<B>;
    fn load_record(self, record: Self::Record) -> Self;
    fn into_record(self) -> Self::Record;

    // Provided methods
    fn devices(&self) -> Vec<<B as Backend>::Device> { ... }
    fn no_grad(self) -> Self { ... }
    fn num_params(&self) -> usize { ... }
    fn save_file<FR, PB>(
        self,
        file_path: PB,
        recorder: &FR
    ) -> Result<(), RecorderError>
       where FR: FileRecorder<B>,
             PB: Into<PathBuf> { ... }
    fn load_file<FR, PB>(
        self,
        file_path: PB,
        recorder: &FR,
        device: &<B as Backend>::Device
    ) -> Result<Self, RecorderError>
       where FR: FileRecorder<B>,
             PB: Into<PathBuf> { ... }
}

Expand description

Trait for all neural network modules.

Modules should be created using the derive attribute. This will make your module trainable, savable and loadable via state and load.

§Example

A module should have a backend defined as a generic parameter B. This will be used by the derive attribute to generate the code necessary to optimize and train the module on any backend.

// Not necessary when using the burn crate directly.
use burn_core as burn;

use burn::{
    nn,
    module::Module,
    tensor::Tensor,
    tensor::backend::Backend,
};

#[derive(Module, Debug)]
struct MyModule<B: Backend> {
  my_param: nn::Linear<B>,
  my_other_field: usize,
}

Required Associated Types§

type Record: Record

Type to save and load the module.

Required Methods§

fn collect_devices( &self, devices: Vec<::Device> ) -> Vec<::Device>

Return all the devices found in the underneath module tree added to the given vector without duplicates.

fn fork(self, device: &::Device) -> Self

Fork the module and all of its sub-modules to the given device.

§Notes

This is similar to to_device, but it ensures the module will have its own autodiff graph.

fn to_device(self, device: &::Device) -> Self

Move the module and all of its sub-modules to the given device.

§Warnings

The device operations will be registered in the autodiff graph. Therefore, be sure to call backward only one time even if you have the same module on multiple devices. If you want to call backward multiple times, look into using fork instead.

fn visit<Visitor>(&self, visitor: &mut Visitor)
where Visitor: ModuleVisitor,

Visit each tensor parameter in the module with a visitor.

fn map<Mapper>(self, mapper: &mut Mapper) -> Self
where Mapper: ModuleMapper,

Map each tensor parameter in the module with a mapper.

fn load_record(self, record: Self::Record) -> Self

Load the module state from a record.

fn into_record(self) -> Self::Record

Convert the module into a record containing the state.

Provided Methods§

fn devices(&self) -> Vec<::Device>

Return all the devices found in the underneath module tree without duplicates.

fn no_grad(self) -> Self

Each tensor in the module tree will not require grad.

§Warnings

This should not be used for inference, use valid when using AD modules. This is mostly useful when performing partial finetuning, which is updating only a small fraction of the parameters instead of finetuning all of them.