Trait Module

Source

pub trait Module:
    Clone
    + Send
    + Debug {
Show 21 methods    // Required methods
    fn collect_devices(&self, devices: Devices) -> Devices;
    fn fork(self, device: &Device) -> Self;
    fn to_device(self, device: &Device) -> Self;
    fn visit<Visitor: ModuleVisitor>(&self, visitor: &mut Visitor);
    fn map<Mapper: ModuleMapper>(self, mapper: &mut Mapper) -> Self;

    // Provided methods
    fn devices(&self) -> Devices { ... }
    fn no_grad(self) -> Self { ... }
    fn freeze_group(self, group: ParamGroup) -> Self { ... }
    fn unfreeze_group(self, group: ParamGroup) -> Self { ... }
    fn train(self) -> Self
       where Self: AutodiffModule { ... }
    fn num_params(&self) -> usize { ... }
    fn quantize_weights(self, quantizer: &mut Quantizer) -> Self { ... }
    fn quantize_weights_group(
        self,
        quantizer: &mut Quantizer,
        group: ParamGroup,
    ) -> Self { ... }
    fn apply_lora(self, config: LoraConfig) -> Self
       where Self: Sized { ... }
    fn apply_qlora(self, config: LoraConfig, quantizer: Quantizer) -> Self
       where Self: Sized { ... }
    fn into_record(self) -> ModuleRecord
       where Self: Sized { ... }
    fn try_load_record(self, record: ModuleRecord) -> Result<Self, RecordError>
       where Self: Sized { ... }
    fn load_record(self, record: ModuleRecord) -> Self
       where Self: Sized { ... }
    fn save_file<P: AsRef<Path>>(self, path: P) -> Result<(), RecordError>
       where Self: Sized { ... }
    fn load_file<P: AsRef<Path>>(self, path: P) -> Self
       where Self: Sized { ... }
    fn try_load_file<P: AsRef<Path>>(self, path: P) -> Result<Self, RecordError>
       where Self: Sized { ... }
}

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

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

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

#[derive(Module, Debug)]
struct MyModule {
  my_param: Linear,
  my_other_field: usize,
}

Required Methods§

Source

fn collect_devices(&self, devices: Devices) -> Devices

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

Source

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 output module on the new device will have its own autodiff graph.

Source

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

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

§Warnings

The operation supports autodiff and it will be registered when activated. However, this may not be what you want. The output model will be an intermediary model, meaning that you can’t optimize it with gradient descent. If you want to optimize the output network on the target device, use fork instead.

Source

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

Visit each tensor parameter in the module with a visitor.

Source

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

Map each tensor parameter in the module with a mapper.

Provided Methods§

Source

fn devices(&self) -> Devices

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

Source

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.

Source

fn freeze_group(self, group: ParamGroup) -> Self

Set require_grad to false for every parameter in the given group, leaving the rest of the module untouched.

This is the group-scoped counterpart to no_grad: where no_grad freezes the whole module tree, freeze_group freezes only the parameters matched by group.

§Warnings

Like no_grad, this should not be used for inference; use valid with AD modules instead.

Source

fn unfreeze_group(self, group: ParamGroup) -> Self

Set require_grad to true for every parameter in the given group, leaving the rest of the module untouched.

The inverse of freeze_group: it re-enables gradient tracking for the parameters matched by group, e.g. to unfreeze a previously frozen module.

Source

fn train(self) -> Self
where Self: AutodiffModule,

Move the module and all of its sub-modules to the autodiff backend.

§Notes

Only plain modules (not already on an autodiff backend) can be moved.
Calling train() on a module that is already on an autodiff backend will result in a type error, because the module’s inner backend does not match.

Source

fn num_params(&self) -> usize

Get the number of parameters the module has, including all of its sub-modules.

Source

fn quantize_weights(self, quantizer: &mut Quantizer) -> Self

Quantize the weights of the module.

Source

fn quantize_weights_group( self, quantizer: &mut Quantizer, group: ParamGroup, ) -> Self

Quantize the weights of the given parameter group.

Source

fn apply_lora(self, config: LoraConfig) -> Self
where Self: Sized,

Attach LoRA adapters to the module’s 2-D weights, freezing the base weights.

The same module keeps working without any code changes; adapted weights now produce base + scale * (a @ b), and only the adapter factors are trainable.

Source

fn apply_qlora(self, config: LoraConfig, quantizer: Quantizer) -> Self
where Self: Sized,

Apply QLoRA to the module: quantize the (frozen) base weights and attach trainable LoRA adapters to 2-D weights.

Source

fn into_record(self) -> ModuleRecord
where Self: Sized,

Collect this module’s parameters into a ModuleRecord.

The record can be saved to a burnpack file or byte buffer and applied back with load_record.

Source

fn try_load_record(self, record: ModuleRecord) -> Result<Self, RecordError>
where Self: Sized,

Apply a ModuleRecord to this module, returning the loaded module.

Honors the record’s DTypePolicy, validate, and allow_partial settings.

Source

fn load_record(self, record: ModuleRecord) -> Self
where Self: Sized,

Apply a ModuleRecord to this module, consuming and returning it.

Panics if validation fails; use try_load_record for the fallible variant.

Source

fn save_file<P: AsRef<Path>>(self, path: P) -> Result<(), RecordError>
where Self: Sized,

Available on crate feature std only.

Save this module’s parameters to a burnpack file on disk.

Convenience for into_record followed by ModuleRecord::save. For non-default load behavior (dtype policy, partial loading, validation), go through the record directly.

Source

fn load_file<P: AsRef<Path>>(self, path: P) -> Self
where Self: Sized,

Available on crate feature std only.

Load this module’s parameters from a burnpack file on disk, returning the loaded module.

Uses the default load behavior. Panics on I/O or validation errors; use try_load_file for the fallible variant, or go through ModuleRecord to configure dtype policy, partial loading or validation.

Source

fn try_load_file<P: AsRef<Path>>(self, path: P) -> Result<Self, RecordError>
where Self: Sized,

Available on crate feature std only.

Fallible variant of load_file.

Reads the record from path with ModuleRecord::load and applies it through try_load_record.

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Trait Module Copy item path

§Example

Required Methods§

fn collect_devices(&self, devices: Devices) -> Devices

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

§Notes

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

§Warnings

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

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

Provided Methods§

fn devices(&self) -> Devices

fn no_grad(self) -> Self

§Warnings

fn freeze_group(self, group: ParamGroup) -> Self

§Warnings

fn unfreeze_group(self, group: ParamGroup) -> Self

fn train(self) -> Selfwhere Self: AutodiffModule,

§Notes

fn num_params(&self) -> usize

fn quantize_weights(self, quantizer: &mut Quantizer) -> Self

fn quantize_weights_group( self, quantizer: &mut Quantizer, group: ParamGroup, ) -> Self

fn apply_lora(self, config: LoraConfig) -> Selfwhere Self: Sized,

fn apply_qlora(self, config: LoraConfig, quantizer: Quantizer) -> Selfwhere Self: Sized,

fn into_record(self) -> ModuleRecordwhere Self: Sized,

fn try_load_record(self, record: ModuleRecord) -> Result<Self, RecordError>where Self: Sized,

fn load_record(self, record: ModuleRecord) -> Selfwhere Self: Sized,

fn save_file<P: AsRef<Path>>(self, path: P) -> Result<(), RecordError>where Self: Sized,

fn load_file<P: AsRef<Path>>(self, path: P) -> Selfwhere Self: Sized,

fn try_load_file<P: AsRef<Path>>(self, path: P) -> Result<Self, RecordError>where Self: Sized,

Dyn Compatibility§

Implementations on Foreign Types§

impl Module for String

fn visit<V: ModuleVisitor>(&self, _visitor: &mut V)

fn map<M: ModuleMapper>(self, _mapper: &mut M) -> Self

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

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

fn collect_devices(&self, devices: Devices) -> Devices

impl Module for bool

fn visit<V: ModuleVisitor>(&self, _visitor: &mut V)

fn map<M: ModuleMapper>(self, _mapper: &mut M) -> Self

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

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

fn collect_devices(&self, devices: Devices) -> Devices

impl Module for f32

fn visit<V: ModuleVisitor>(&self, _visitor: &mut V)

fn map<M: ModuleMapper>(self, _mapper: &mut M) -> Self

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

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

fn collect_devices(&self, devices: Devices) -> Devices

impl Module for f64

fn visit<V: ModuleVisitor>(&self, _visitor: &mut V)

fn map<M: ModuleMapper>(self, _mapper: &mut M) -> Self

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

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

fn collect_devices(&self, devices: Devices) -> Devices

impl Module for i8

fn visit<V: ModuleVisitor>(&self, _visitor: &mut V)

fn map<M: ModuleMapper>(self, _mapper: &mut M) -> Self

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

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

fn collect_devices(&self, devices: Devices) -> Devices

impl Module for i16

fn visit<V: ModuleVisitor>(&self, _visitor: &mut V)

fn map<M: ModuleMapper>(self, _mapper: &mut M) -> Self

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

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

fn collect_devices(&self, devices: Devices) -> Devices

impl Module for i32

fn visit<V: ModuleVisitor>(&self, _visitor: &mut V)

fn map<M: ModuleMapper>(self, _mapper: &mut M) -> Self

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

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

fn collect_devices(&self, devices: Devices) -> Devices

impl Module for i64

fn visit<V: ModuleVisitor>(&self, _visitor: &mut V)

fn map<M: ModuleMapper>(self, _mapper: &mut M) -> Self

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

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

fn collect_devices(&self, devices: Devices) -> Devices

Trait Module

fn train(self) -> Self
where Self: AutodiffModule,

fn apply_lora(self, config: LoraConfig) -> Self
where Self: Sized,

fn apply_qlora(self, config: LoraConfig, quantizer: Quantizer) -> Self
where Self: Sized,

fn into_record(self) -> ModuleRecord
where Self: Sized,

fn try_load_record(self, record: ModuleRecord) -> Result<Self, RecordError>
where Self: Sized,

fn load_record(self, record: ModuleRecord) -> Self
where Self: Sized,

fn save_file<P: AsRef<Path>>(self, path: P) -> Result<(), RecordError>
where Self: Sized,

fn load_file<P: AsRef<Path>>(self, path: P) -> Self
where Self: Sized,

fn try_load_file<P: AsRef<Path>>(self, path: P) -> Result<Self, RecordError>
where Self: Sized,