espnet3.components.modeling.lightning_module.ESPnetLightningModule
espnet3.components.modeling.lightning_module.ESPnetLightningModule
class espnet3.components.modeling.lightning_module.ESPnetLightningModule(model, config)
Bases: LightningModule
ESPnet3 LightningModule wrapper for model training and data integration.
This module follows Lightning best practices by defining the training/validation steps, optimizer/scheduler configuration, and dataloader hooks on the LightningModule itself. It also integrates ESPnet-specific dataset handling, distributed NaN/Inf loss skipping, and statistics collection.
model
The main ESPnet model.
- Type: torch.nn.Module
config
Training configuration for model and data setup.
- Type: DictConfig
train_dataset
Training dataset organized by the ESPnet data organizer.
valid_dataset
Validation dataset.
collate_fn
Collation function used in DataLoader.
- Type: Callable
is_espnet_sampler
Whether the model uses ESPnet’s custom sampler.
- Type: bool
NOTE
This class assumes the use of a DataOrganizer-compatible dataset config. The data_organizer is instantiated temporarily to access train and valid datasets, but is not retained as an attribute since it is no longer needed after extraction.
Initialize the ESPnet LightningModule wrapper.
collect_stats()
Collect training and validation statistics using ESPnet’s collect_stats.
Requires config.stats_dir to be defined. Saves stats under this directory.
- Raises:AssertionError – If config.stats_dir is not provided.
config
Configure optimizers and schedulers for training.
This method supports two modes of configuration.
- Single Optimizer + Scheduler: Use when the entire model is trained with a single optimizer.
```yaml optim:
_target_: torch.optim.Adam lr: 0.001
scheduler: : _target_: torch.optim.lr_scheduler.StepLR step_size: 10
- Multiple Optimizers + Schedulers: Use when training different parts of the model with different optimizers. Each optimizer block must contain both a nested optim config and a params key indicating a substring to match parameter names.
```yaml optims:
- optim: : _target_: torch.optim.Adam lr: 0.001
params: encoder
- optim: : _target_: torch.optim.SGD lr: 0.01
params: decoder
schedulers: : - scheduler: : _target_: torch.optim.lr_scheduler.StepLR step_size: 10
- scheduler: : _target_: torch.optim.lr_scheduler.ReduceLROnPlateau patience: 2
```Notes
Only one of optim or optims may be specified. Mixing is not allowed.
Likewise, scheduler and schedulers must not be used together.
When using optims, each params must uniquely match a subset of : trainable parameters.
It is an error if: : * A trainable parameter is assigned to multiple optimizers : (overlapping params)
- A trainable parameter is not assigned to any optimizer
- Any optimizer block is missing params or nested optim
Returns: A dictionary with keys “optimizer” and “lr_scheduler” : for PyTorch Lightning.
Return type: dict
Raises:
- AssertionError – If configuration rules are violated.
- ValueError – If neither optimizer configuration is provided.
load_state_dict(state_dict, strict=True)
Load state dict into the model.
state_dict(*args, **kwargs)
Return the state dict of the model.
train_dataloader()
Build the training DataLoader using ESPnet’s DataLoaderBuilder.
- Returns: The training DataLoader.
- Return type: DataLoader
training_step(batch, batch_idx)
Training step logic.
val_dataloader()
Build the validation DataLoader using ESPnet’s DataLoaderBuilder.
- Returns: The validation DataLoader.
- Return type: DataLoader
validation_step(batch, batch_idx)
Validate step.