Trait Module

Source

pub trait Module:
    Clone
    + Send
    + Debug {
Show 28 methods    // Required methods
    fn collect_devices(&self, devices: Devices) -> Devices;
    fn fork(self, device: &Device) -> Self;
    fn to_device(self, device: &Device) -> Self;
    fn train(self) -> Self;
    fn valid(&self) -> 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 set_require_grad(self, require_grad: bool) -> Self { ... }
    fn set_require_grad_group(
        self,
        group: ParamGroup,
        require_grad: bool,
    ) -> Self { ... }
    fn no_grad(self) -> Self { ... }
    fn freeze(self) -> Self { ... }
    fn unfreeze(self) -> Self { ... }
    fn freeze_group(self, group: ParamGroup) -> Self { ... }
    fn unfreeze_group(self, group: ParamGroup) -> Self { ... }
    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_reparameterization<R>(self, reparameterizer: R) -> Self
       where Self: Sized,
             R: Reparameterizer { ... }
    fn apply_lora(self, lora: Lora) -> Self
       where Self: Sized { ... }
    fn apply_qlora(self, qlora: QLora) -> Self
       where Self: Sized { ... }
    fn into_record(self) -> ModuleRecord
       where Self: Sized { ... }
    fn into_record_group(self, group: ParamGroup) -> 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 into_record and load_record.

§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,
}

Training and validation use the same module type. This trait provides state transitions, not a guarantee that a value currently has autodiff enabled. Inspect individual parameter tensors with Tensor::is_autodiff and Tensor::is_require_grad; contexts can be mixed.

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. Both transfers preserve the source tensors’ autodiff association and checkpointing strategy. The destination’s autodiff defaults do not enable training; use train first when starting from a validation module.

A parameter not initialized yet initializes on the destination, unless a clone shares it, in which case it initializes where it is and is then copied.

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.

A parameter not initialized yet initializes on the destination, unless a clone shares it, in which case it initializes where it is and is then moved.

Source

fn train(self) -> Self

Enables autodiff and restores configured training state throughout the module.

§Gradient and training state

This is the supported transition back to training after valid. It applies the latest gradient-tracking settings and training-flag values. Mappings, device moves, forks and explicit state changes on a validation module preserve or update what this method applies.

Tensor-bearing modules require the autodiff feature to enable autodiff. This does not undo explicit freezing or reconstruct state discarded by valid.

Source

fn valid(&self) -> Self

Returns a validation snapshot without autodiff or active training flags.

§Gradient and training state

Tensor gradient requirements and module-owned training flags are disabled in the returned value while their configured state is retained internally. Use Module::train to apply that state. Mappings, device moves and forks preserve both the disabled effective validation state and what train will restore. Calling valid when the module’s tensors are already on a plain device still disables its training flags, so the operation is idempotent but not necessarily a structural no-op.

The returned value is an inference snapshot: parameter reparameterizations are folded into their values, and tensor checkpointing strategies are removed with autodiff. Calling Module::train on that snapshot does not reconstruct those reparameterizations or restore the previous checkpointing strategies. Keep the original module when continuing training after validation.

Source

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

Visit each parameter and module-owned control value with a visitor.

Source

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

Map each parameter and module-owned control value with a mapper.

Provided Methods§

Source

fn devices(&self) -> Devices

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

Source

fn set_require_grad(self, require_grad: bool) -> Self

Set whether every floating-point tensor parameter in the module tree requires gradients.

Module-owned control flags are left unchanged. Use freeze to disable both gradients and training behavior. On a backend without autodiff, the enabled setting is preserved and takes effect when train transitions the module back to training. Setting this to true overrides selective gradient configurations such as the frozen dense base weights established by apply_lora.

Source

fn set_require_grad_group(self, group: ParamGroup, require_grad: bool) -> Self

Set whether matched floating-point tensor parameters require gradients.

Only floating-point tensor parameters within group are affected. Other parameter values matched by the group, including module-owned control flags, are ignored; unmatched tensor parameters are also left unchanged. This means a group created with ParamGroup::ids_from_module can select an entire subtree while this method changes only its tensor gradient state. On a backend without autodiff, the enabled setting is preserved and takes effect through train.

Use freeze_group or unfreeze_group when matched training flags should be changed together with tensor gradients.

Source

fn no_grad(self) -> Self

Disable gradients for every floating-point tensor parameter in the module tree.

This is equivalent to set_require_grad(false). Module-owned control flags are left unchanged, so dropout, noise, randomized activations and batch-normalization running statistics retain their current behavior. Use freeze to disable both gradients and training behavior.

§Gradient and training state

This persistently disables gradient tracking, including on a backend without autodiff. A later train keeps gradients disabled; it does not undo no_grad.

§Warnings

Use valid for inference. no_grad is useful for partial fine-tuning when tensor gradients should be disabled while layer training behavior remains active. For example, matched batch-normalization layers continue updating their running statistics, and dropout remains enabled. Use freeze when that module-owned training behavior should also stop.

Source

fn freeze(self) -> Self

Freeze the whole module tree.

Every floating-point tensor parameter stops requiring gradients, and every module-owned training flag is disabled. This applies both to parameterless stochastic layers such as dropout, noise and randomized activations, and to stateful layers such as batch normalization.

§Gradient and training state

Gradient tracking and training flags are persistently disabled. On a backend without autodiff, flags are disabled immediately, and train keeps tensors untracked rather than undoing the freeze.

§Warnings

Use valid for inference. freeze is intended for partial finetuning where a module remains on the training device but should not participate in training.

Source

fn unfreeze(self) -> Self

Unfreeze the whole module tree.

Every floating-point tensor parameter is configured to require gradients, and every module-owned training flag is set to enabled rather than restored to a prior value. This overwrites selective gradient configurations, including frozen dense base weights established by apply_lora.

§Gradient and training state

On a backend without autodiff, training flags are enabled immediately while gradient tracking takes effect when train transitions the module back to the autodiff backend.

Source

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

Freeze every module-owned value in the given group, leaving the rest of the module untouched.

This is the group-scoped counterpart to freeze: where freeze freezes the whole module tree, freeze_group clears gradient requirements on matched tensor parameters and disables matched training flags.

§Gradient and training state

Gradient tracking and training flags are persistently disabled for matched values. On a backend without autodiff, matched flags are disabled immediately, and train keeps matched tensors untracked rather than undoing the group freeze.

§Warnings

Like freeze, this should not be used for inference; use valid instead.

Source

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

Unfreeze every module-owned value in the given group, leaving the rest of the module untouched.

The inverse of freeze_group: matched tensor parameters are configured to require gradients, and matched training flags are set to enabled rather than restored to prior values.

§Gradient and training state

On a backend without autodiff, matched training flags are enabled immediately while gradient tracking for matched tensors takes effect when train transitions the module back to the autodiff backend.

Source

fn num_params(&self) -> usize

Get the number of parameters the module has, including all of its sub-modules, without initializing the ones not initialized yet.

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_reparameterization<R>(self, reparameterizer: R) -> Self
where Self: Sized, R: Reparameterizer,

Attach reparameterizations using the given Reparameterizer.

Every floating-point parameter is passed to the reparameterizer along with its module path. The reparameterizer prepares its structural base and optionally creates the state used by Param::val. Lora is a built-in example.

§Limitations

Nested reparameterizations are not supported. This method should only be called on modules that don’t already contain reparameterized parameters.

Source

fn apply_lora(self, lora: Lora) -> Self
where Self: Sized,

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

The same module keeps working without any code changes; adapted weights now produce base + scale * (a @ b), and only the adapter factors are trainable. Module-owned control flags are preserved, so dropout and batch normalization keep their current training behavior. Call freeze before this method if the base module’s control behavior and running statistics should also be frozen.

Source

fn apply_qlora(self, qlora: QLora) -> Self
where Self: Sized,

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

Module-owned control flags are preserved. Call freeze before this method if the base module’s control behavior and running statistics should also be frozen.

Source

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

Collect this module’s tensor parameters into a ModuleRecord.

The record can be saved to a burnpack file or byte buffer and applied back with load_record. Module-owned control values such as Flag are runtime configuration and are not recorded; loading preserves their state and identity from the destination module.

Source

fn into_record_group(self, group: ParamGroup) -> ModuleRecord
where Self: Sized,

Collect the tensor parameters group names into a ModuleRecord.

The record of a part of the module rather than all of it — what a run that trained a group writes when the rest of the module is the checkpoint it started from, and what load_record applies back over that checkpoint (with allow_partial, since the record holds nothing for the parameters outside the group).

A parameter the group does not match is skipped before its data is read, so this never materializes the rest of the module.

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 train(self) -> Self

§Gradient and training state

fn valid(&self) -> Self

§Gradient and training state

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

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

Provided Methods§

fn devices(&self) -> Devices

fn set_require_grad(self, require_grad: bool) -> Self

fn set_require_grad_group(self, group: ParamGroup, require_grad: bool) -> Self

fn no_grad(self) -> Self

§Gradient and training state

§Warnings

fn freeze(self) -> Self

§Gradient and training state

§Warnings

fn unfreeze(self) -> Self

§Gradient and training state

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

§Gradient and training state

§Warnings

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

§Gradient and training state

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_reparameterization<R>(self, reparameterizer: R) -> Selfwhere Self: Sized, R: Reparameterizer,

§Limitations

fn apply_lora(self, lora: Lora) -> Selfwhere Self: Sized,

fn apply_qlora(self, qlora: QLora) -> Selfwhere Self: Sized,

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

fn into_record_group(self, group: ParamGroup) -> 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

fn valid(&self) -> Self

fn train(self) -> Self

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

fn valid(&self) -> Self

fn train(self) -> Self

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

fn valid(&self) -> Self

fn train(self) -> Self

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

fn valid(&self) -> Self

fn train(self) -> Self

impl Module for i8

Trait Module

fn apply_reparameterization<R>(self, reparameterizer: R) -> Self
where Self: Sized, R: Reparameterizer,

fn apply_lora(self, lora: Lora) -> Self
where Self: Sized,

fn apply_qlora(self, qlora: QLora) -> Self
where Self: Sized,

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

fn into_record_group(self, group: ParamGroup) -> 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,