PyTorch Extension Points

This reference describes the contracts to use when extending Hyperion’s PyTorch stack. It is intentionally selective: the public base classes, factories, and data interfaces below are the supported places to integrate new models and training behavior. For a runnable waveform x-vector workflow, see Train an X-Vector Directly from Waveforms.

Layering and ownership

Hyperion separates a reusable network from the task that trains it:

layers contain primitive operations;
layer_blocks compose primitives into reusable modules;
narchs contain network architectures and report their tensor shapes;
models own task-specific forward, loss, and embedding behavior;
data produces batches and sampling plans; and
trainers own the optimization loop, checkpoints, logging, AMP, and distributed execution.

Add a reusable encoder or decoder to narchs. Add a model only when it defines task-level behavior. Do not put checkpointing or distributed-launch logic in either: that responsibility belongs to a trainer.

Neural architectures

NetArch is the architecture-level base class. Its shape methods describe tensor shapes including the batch axis. Implement them accurately: model wrappers and configuration validation depend on the reported channel and time dimensions. in_context expresses any frame context the architecture needs.

The maintained architecture families include TDNN, ResNet/ResNet1d, Conformer, Transformer, ConvNeXt, EfficientNet, and SpineNet. Select one that matches the input representation rather than copying its implementation into a task model. Architectures that are specific to codec/DAC or transducer systems are experimental.

Task models and waveform wrappers

class hyperion.torch.HyperTorchModel(*args: Any, **kwargs: Any)[source]

Base class for PyTorch models and neural network architectures.

registry: Dict[str, Type[HyperTorchModel]] = {'TorchModel': <class 'hyperion.torch.hyper_torch_model.TorchModel'>}

classmethod __init_subclass__(**kwargs: Any) → None[source]

Register subclasses by class name for dynamic loading.

Parameters:: **kwargs – Additional subclass initialization options.

__init__(bias_weight_decay: float | None = None)[source]

Initialize model-level training controls.

Parameters:: bias_weight_decay – Optional decay value for bias/1D parameter group.

get_config(no_class_name: bool = False) → Dict[str, Any][source]

Return a serializable configuration dictionary.

Parameters:: no_class_name – If True, omit class_name from the config.
Returns:: Configuration dictionary used to reconstruct the model.

copy() → HyperTorchModel[source]

Return a deep copy of this model.

Returns:: A deep-copied model instance.

clone() → HyperTorchModel[source]

Return a deep copy of this model (alias of copy).

Returns:: A deep-copied model instance.

trainable_parameters(recurse: bool = True) → Iterator[torch.nn.Parameter][source]

Yield parameters with requires_grad=True.

Parameters:: recurse – Whether to recurse into submodules.
Returns:: Iterator of trainable parameters.

non_trainable_parameters(recurse: bool = True) → Iterator[torch.nn.Parameter][source]

Yield parameters with requires_grad=False.

Parameters:: recurse – Whether to recurse into submodules.
Returns:: Iterator of non-trainable parameters.

trainable_named_parameters(recurse: bool = True) → Iterator[Tuple[str, torch.nn.Parameter]][source]

Yield (name, parameter) pairs for trainable parameters.

Parameters:: recurse – Whether to recurse into submodules.
Returns:: Iterator of named trainable parameters.

non_trainable_named_parameters(recurse: bool = True) → Iterator[Tuple[str, torch.nn.Parameter]][source]

Yield (name, parameter) pairs for non-trainable parameters.

Parameters:: recurse – Whether to recurse into submodules.
Returns:: Iterator of named non-trainable parameters.

parameter_summary(verbose: bool = False) → Tuple[int, int, int, int, int][source]

Return parameter and buffer counts.

Parameters:: verbose – If True, log a summary line.
Returns:: Tuple (total, trainable, non_trainable_plus_buffers, non_trainable, buffers).

print_parameter_list() → None[source]: Log names of trainable, non-trainable, and buffer tensors.

has_param_groups() → bool[source]

Return whether model exposes custom optimizer parameter groups.

Returns:: True if bias_weight_decay is enabled.

trainable_param_groups() → List[Dict[str, Any]][source]

Build optimizer parameter groups for trainable parameters.

Returns:: Optimizer parameter-group dictionaries.

freeze() → None[source]: Disable gradients for all parameters.

unfreeze() → None[source]: Enable gradients for all parameters.

has_batchnorms() → bool[source]

Return True if the model contains any batch-normalization layer.

Returns:: True when any batchnorm module is found.

change_dropouts(dropout_rate: float) → None[source]

Set dropout probability on dropout and RNN modules.

Parameters:: dropout_rate – New dropout rate.

property train_mode: str

full or frozen.

Returns:: Active train mode string.
Type:: Current train mode

set_train_mode(mode: str) → None[source]

Switch model between full-train and frozen-parameter modes.

Parameters:: mode – Target mode (full or frozen).

_train(train_mode: str) → None[source]

Apply PyTorch train/eval state for a custom train mode.

Parameters:: train_mode – Mode to apply (full or frozen).

train(mode: bool = True) → HyperTorchModel[source]

Override nn.Module.train to honor self.train_mode.

Parameters:: mode – If False, force eval mode.
Returns:: self.

static valid_train_modes() → List[str][source]

Return the list of supported train modes.

Returns:: Supported train-mode names.

save(file_path: str | Path) → None[source]

Save model config and state dictionary to disk.

Parameters:: file_path – Destination checkpoint path.

Resolve config/state_dict from args or a checkpoint file.

Parameters:

file_path – Optional checkpoint path.
cfg – Optional pre-loaded model config.
state_dict – Optional pre-loaded state dict.

Returns:

Tuple (cfg, state_dict) with class_name removed from config.

classmethod load(file_path: str | Path | None = None, cfg: Dict[str, Any] | None = None, state_dict: Dict[str, torch.Tensor] | None = None) → HyperTorchModel[source]

Instantiate model from config and optionally load weights.

Parameters:

file_path – Optional checkpoint path.
cfg – Optional model config.
state_dict – Optional model weights.

Returns:

Instantiated model.

get_reg_loss() → int[source]

Return regularization loss contribution for this model.

Returns:: Regularization loss term.

get_loss() → int[source]

Return auxiliary loss contribution for this model.

Returns:: Auxiliary loss term.

property device: torch.device

Return unique device shared by parameters and buffers.

Returns:: Device where all parameters and buffers reside.

static _remove_module_prefix(state_dict: Dict[str, torch.Tensor]) → Dict[str, torch.Tensor][source]

Remove leading module. prefixes from state-dict keys.

Parameters:: state_dict – State dict to normalize.
Returns:: Normalized state dict.

static _fix_xvector_cfg(cfg: Dict[str, Any]) → Dict[str, Any][source]

Normalize legacy XVector config keys to current names.

Parameters:: cfg – Model config dictionary.
Returns:: Updated config dictionary.

static _fix_hf_wav2xvector(cfg: Dict[str, Any], state_dict: Dict[str, torch.Tensor]) → Tuple[Dict[str, Any], Dict[str, torch.Tensor]][source]

Migrate legacy HF wav2xvector fusion config and checkpoint layout.

Parameters:

cfg – Model config dictionary.
state_dict – Model checkpoint weights.

Returns:

Updated (cfg, state_dict).

static _fix_resnet_qvector_cfg(cfg: Dict[str, Any]) → Dict[str, Any][source]

Drop deprecated ResNetQVector config keys.

Parameters:: cfg – Model config dictionary.
Returns:: Updated config dictionary.

static _fix_model_compatibility(class_obj: Type[HyperTorchModel], cfg: Dict[str, Any], state_dict: Dict[str, torch.Tensor]) → Tuple[Dict[str, Any], Dict[str, torch.Tensor]][source]

Apply compatibility fixes for deprecated model formats.

Parameters:

class_obj – Model class to instantiate.
cfg – Configuration dictionary.
state_dict – Serialized model weights.

Returns:

Updated (cfg, state_dict).

static _is_hf_path(file_path: Path) → bool[source]

Return True for HF-style org/repo/file paths.

Parameters:: file_path – Path to validate.
Returns:: True if path shape matches expected HF format.

static _get_from_hf(file_path: Path, cache_dir: str | Path = None, local_dir: str | Path = None) → str[source]

Download a file from Hugging Face Hub and return its local path.

Parameters:

file_path – HF-style path org/repo/file.
cache_dir – Optional HF cache directory.
local_dir – Optional local download directory.

Returns:

Local filesystem path to downloaded file.

static _try_to_get_from_hf(file_path: Path, cache_dir: str | Path = None, local_dir: str | Path = None) → Path[source]

Resolve local file path, downloading from HF Hub when needed.

Parameters:

file_path – Local path or prefixed HF path.
cache_dir – Optional HF cache directory.
local_dir – Optional local download directory.

Returns:

Resolved local path.

static _bootstrap_registry() → None[source]

Import common torch subpackages so subclasses register themselves.

HyperTorchModel.registry is populated when subclass definitions are imported. CLI scripts often import only HyperTorchModel and call auto_load, so we need a lazy import step before class lookup.

static _find_module_for_class_name(class_name: str) → str | None[source]

Find module path for class name by scanning torch source files.

Parameters:: class_name – Target class name.
Returns:: Dotted module path if found, otherwise None.

Load model from checkpoint with dynamic class resolution.

Parameters:

file_path – Local path or HF-style path.
model_name – State-dict prefix; defaults to model.
extra_objs – Optional mapping from class names to model classes.
map_location – torch.load map location.
cache_dir – Optional HF cache directory.
local_dir – Optional local download directory.

Returns:

Loaded model instance.

HyperTorchModel is the serializable base for trainable task models. A subclass must retain a JSON-friendly configuration and use the inherited save/load mechanism so that auto_load can recreate it from an artifact. The model’s forward contract is task-specific; document its input tensor layout, length/mask semantics, target fields, and returned values alongside the model.

Wav2XVector accepts waveforms shaped (batch, num_samples). It performs toolkit acoustic feature extraction before calling the x-vector backend. HFWav2XVector follows the same task contract but uses a Hugging Face feature extractor and fuses selected hidden layers. Its pretrained checkpoint, cache, and fine-tuning policy should be explicit in the model configuration; see Train a Pretrained Wav2Vec2 X-Vector.

Datasets and samplers

AudioDataset is the current dataset interface for waveform training. It loads a HyperDataset and can return class labels, tokenized attributes, extra metadata, and augmentations. LegacyAudioDataset is kept for existing CSV-manifest x-vector workflows; prefer AudioDataset for new integrations unless the maintained command you are extending requires the legacy batch format.

Samplers emit batch index lists and are responsible for reproducible rank-aware ordering. Call set_epoch on resume and at every epoch boundary. Use max_batch_length to bound padded waveform cost; it is generally a more reliable memory control than only setting a fixed batch size. The sampler factory is the public configuration boundary for the sequence sampler family.

Training and optimization

TorchTrainerBase owns generic checkpoint and launch policy. The x-vector trainers define the speaker-classification batch and loss convention; XVectorTrainerFromWav additionally applies an acoustic feature extractor. Choose an existing trainer before creating one: custom models normally only need to satisfy the selected trainer’s model and batch contract. See Run Resumable, Mixed-Precision, and Distributed Training for resume, AMP, and DDP.

class hyperion.torch.lr_schedulers.factory.LRSchedulerFactory[source]

Factory for creating configured learning-rate schedulers.

static create(optimizer: torch.optim.Optimizer, lrsch_type: str, decay_rate: float = 0.01, decay_steps: int = 100, power: float = 0.5, hold_steps: int = 10, t: int = 10, t_mul: int = 1, warm_restarts: bool = False, gamma: float = 1, monitor: str = 'val_loss', mode: str = 'min', factor: float = 0.1, patience: int = 10, threshold: float = 0.0001, threshold_mode: str = 'rel', cooldown: int = 0, eps: float = 1e-08, min_lr: float | Sequence[float] = 0, warmup_steps: int | None = None, d_model: int | None = None, lr_factor: float = 1, update_lr_on_opt_step: bool = True) → LRScheduler | None[source]

Create a learning-rate scheduler instance.

Parameters:

optimizer – Wrapped optimizer.
lrsch_type – Scheduler type identifier.
decay_rate – Exponential decay factor.
decay_steps – Number of steps associated with one exponential decay.
power – Inverse-power decay exponent.
hold_steps – Steps to hold initial LR before decay.
t – Base cycle length for cyclic schedulers.
t_mul – Cycle-length multiplier after each restart.
warm_restarts – Enable warm restarts for cosine schedule.
gamma – Max-LR multiplier after each restart.
monitor – Metric key for plateau scheduler.
mode – "min" or "max" for plateau scheduler.
factor – LR reduction factor for plateau scheduler.
patience – Patience (epochs) for plateau scheduler.
threshold – Improvement threshold for plateau scheduler.
threshold_mode – "rel" or "abs" threshold semantics.
cooldown – Cooldown epochs for plateau scheduler.
eps – Minimum effective LR change for plateau scheduler.
min_lr – Scalar or per-group lower LR bound.
warmup_steps – Linear warmup duration in optimizer steps. Uses scheduler-specific defaults when omitted.
d_model – Transformer hidden size for Noam schedule.
lr_factor – Scale factor for Noam schedule.
update_lr_on_opt_step – Whether to update LR on optimizer steps.

Returns:

Scheduler instance, or None when lrsch_type == "none".

Raises:

ValueError – If lrsch_type is unknown.

static filter_args(**kwargs: Any) → Dict[str, Any][source]: Filter a kwargs dictionary to args accepted by create().

static add_class_args(parser: ArgumentParser, prefix: str | None = None) → None[source]: Register LR scheduler CLI arguments in an argument parser.

static add_argparse_args(parser: ArgumentParser, prefix: str | None = None) → None: Register LR scheduler CLI arguments in an argument parser.

class hyperion.torch.wd_schedulers.factory.WDSchedulerFactory[source]

Factory for creating configured weight-decay schedulers.

static create(optimizer: torch.optim.Optimizer, wdsch_type: str, initial_wd: float | Sequence[float] = 1e-05, warmup_steps: int = 0, update_wd_on_opt_step: bool = True) → WDScheduler | None[source]

Create a weight-decay scheduler instance.

Parameters:

optimizer – Wrapped optimizer.
wdsch_type – Scheduler type identifier.
initial_wd – Initial weight decay value (scalar or per-group).
warmup_steps – Steps until reaching final weight decay.
update_wd_on_opt_step – If True, update WD on optimizer steps; otherwise update on epoch boundaries.

Returns:

Scheduler instance, or None when wdsch_type == "none".

Raises:

ValueError – If wdsch_type is unknown.

static filter_args(**kwargs: Any) → Dict[str, Any][source]: Filter a kwargs dictionary to args accepted by create().

static add_class_args(parser: ArgumentParser, prefix: str | None = None) → None[source]: Register WD scheduler CLI arguments in an argument parser.

static add_argparse_args(parser: ArgumentParser, prefix: str | None = None) → None: Register WD scheduler CLI arguments in an argument parser.

These factories are the configuration-facing interface for optimizers and schedulers. Add a new optimizer or schedule here only when it is intended for general reuse; a model-specific learning rule belongs with its trainer and should be documented as such.

Stable and experimental families

The general layers, layer blocks, architecture families, waveform x-vector models, datasets/samplers, trainer infrastructure, TPM wrappers, and adversarial components are stable public surfaces. Codec/DAC, VITS anonymization and voice conversion, transducers, and q-vectors remain experimental: their configuration and serialized artifacts may change between releases. Refer to Documentation Policy before building against an experimental family.

PyTorch Extension Points

Layering and ownership

Neural architectures

Task models and waveform wrappers

Datasets and samplers

Training and optimization

Stable and experimental families

See also