Trait Module

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

    // Provided methods
    fn devices(&self) -> Vec<Device> { ... }
    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>(self, path: P) -> Result<(), RecordError>
       where P: AsRef<Path>,
             Self: Sized { ... }
    fn load_file<P>(self, path: P) -> Self
       where P: AsRef<Path>,
             Self: Sized { ... }
    fn try_load_file<P>(self, path: P) -> Result<Self, RecordError>
       where P: AsRef<Path>,
             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§

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.

Device collection should not initialize lazy parameters. Use Param::lazy_device to inspect their devices without accessing tensor values.

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.

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

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

§Warnings

For initialized parameters, the operation supports autodiff and is registered when activated. 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. To train the module on the destination device, use fork. This preserves gradient retention whether the parameters are uninitialized, initialized, or shared with clones.

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 restore tensor checkpointing strategies discarded by valid.

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.

Parameter reparameterizations and their stored bases are preserved, including quantized bases. Their nested modules are also put into validation mode. Use materialize explicitly to fold reparameterizations into weights. Tensor checkpointing strategies are removed with autodiff and are not restored by Module::train. Keep the original module when continuing training after validation.

fn materialize(self) -> Self

Fold parameter reparameterizations into their effective values and remove them.

Each reparameterized parameter is replaced by the value computed from its stored base and reparameterization state. Other parameters and module training flags are unchanged. The resulting parameters retain their base parameter IDs, layout mappings and configured trainability, but are detached from the computation that produced them. Parameters without reparameterizations remain lazy; this does not initialize them.

This consumes the module and cannot be undone on the returned value. Clone the module first to retain its reparameterizations. To prepare merged inference weights, use model.valid().materialize(); materialization alone does not switch to validation mode.

§Memory and precision

LoRA adapters are merged into dense weights. For QLoRA this also replaces the packed base with a dense effective weight, which can substantially increase memory use. Requantization must be requested separately and may change the model’s outputs.

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

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

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

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

Provided Methods§

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

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

Lazy tensor parameters report their initialization device without being initialized. Device equality ignores autodiff and checkpointing settings, so this list is not a summary of the module’s training state.

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.

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.

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.

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.

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.

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.

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.

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.

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

Quantize the weights of the module.

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

Quantize the weights of the given parameter group.

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.

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.

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.

Dense effective weights can still materialize during training, even with fusion enabled. See QLora’s memory limitations before sizing a training run.

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.

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.

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.

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.

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.

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

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.

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

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.