Builders

The builders package provides utilities for constructing neural networks and optimizers from configuration specifications.

Network Builder

class aenet.torch_training.builders.NetworkBuilder(descriptor, device: torch.device, dtype: torch.dtype)[source]

Bases: object

Builds neural network architectures for atomic energy prediction.

Supports both aenet-PyTorch NetAtom (if available) and a fallback implementation using standard PyTorch modules.

Parameters:
  • descriptor (ChebyshevDescriptor) – Descriptor instance providing feature dimension and species info.

  • device (torch.device) – Device for network.

  • dtype (torch.dtype) – Data type for network parameters.

build_network(arch: Dict[str, List[Tuple[int, str]]]) torch.nn.Module[source]

Build NetAtom (preferred) or fallback per-species MLPs.

Parameters:

arch (dict) – Architecture specification per species.

Returns:

Network with attributes: - .functions: ModuleList of per-species Sequential MLPs - .device: device string

Return type:

nn.Module

validate_arch(arch: Dict[str, List[Tuple[int, str]]]) Tuple[List[List[int]], List[List[str]]][source]

Validate architecture and produce per-species hidden sizes and activations.

Parameters:

arch (dict) – {species_symbol: [(nodes, activation), …]} Output layer is implicit.

Returns:

  • hidden_sizes (list of list of int) – Per-species hidden layer sizes.

  • activations (list of list of str) – Per-species activation functions.

Raises:

ValueError – On unsupported activation or missing species.

Optimizer Builder

class aenet.torch_training.builders.OptimizerBuilder(model: torch.nn.Module)[source]

Bases: object

Builds optimizers and learning rate schedulers.

Parameters:

model (nn.Module) – Model whose parameters will be optimized.

build_optimizer(method: TrainingMethod) torch.optim.Optimizer[source]

Build optimizer from training method configuration.

Parameters:

method (TrainingMethod) – Training method configuration (Adam or SGD).

Returns:

Configured optimizer.

Return type:

torch.optim.Optimizer

static build_scheduler(optimizer: torch.optim.Optimizer, use_scheduler: bool = False, scheduler_patience: int = 10, scheduler_factor: float = 0.5, scheduler_min_lr: float = 1e-06) torch.optim.lr_scheduler.ReduceLROnPlateau | None[source]

Build learning rate scheduler.

Parameters:
  • optimizer (torch.optim.Optimizer) – Optimizer to schedule.

  • use_scheduler (bool) – Whether to use a scheduler.

  • scheduler_patience (int) – Number of epochs with no improvement after which LR is reduced.

  • scheduler_factor (float) – Factor by which LR is reduced.

  • scheduler_min_lr (float) – Minimum learning rate.

Returns:

Scheduler if use_scheduler is True, otherwise None.

Return type:

torch.optim.lr_scheduler.ReduceLROnPlateau or None

static get_batch_size(method: TrainingMethod) int[source]

Extract batch size from training method.

Parameters:

method (TrainingMethod) – Training method configuration.

Returns:

Batch size (defaults to 32 if not found).

Return type:

int