pycsamt.ai.training.trainer#

Training loop for EM neural networks.

EMTrainer wraps a PyTorch nn.Module and handles:

  • Mini-batch training with DataLoader

  • Masked MSE loss (ignores NaN padding for variable-depth datasets)

  • ReduceLROnPlateau learning-rate scheduling

  • Early stopping with configurable patience

  • Per-epoch progress display (tqdm if available)

  • History dict for plot_convergence()

Classes

EMTrainer(model, *[, lr, weight_decay, ...])

Training loop manager for EM 1-D inversion networks.
class pycsamt.ai.training.trainer.EMTrainer(model, *, lr=0.001, weight_decay=1e-05, patience=20, min_delta=1e-05, batch_size=256, device='cpu', grad_clip=None, verbose=True)[source]#

Bases: object

Training loop manager for EM 1-D inversion networks.

Parameters:

  • model (nn.Module) – The PyTorch network to train.

  • lr (float) – Initial learning rate. Default 1e-3.

  • weight_decay (float) – L2 regularisation. Default 1e-5.

  • patience (int) – Early-stopping patience (epochs without val-loss improvement).

  • min_delta (float) – Minimum val-loss improvement to reset patience counter.

  • batch_size (int) – Mini-batch size.

  • device (str) – Compute device ('cpu', 'cuda', 'mps').

  • grad_clip (float or None) – If set, clip gradient norms to this value.

  • verbose (bool) – Print per-epoch summary.

Variables:

  • history (dict) – Keys 'train_loss' and 'val_loss' — lists of per-epoch averages. Also 'lr' and 'epoch_time'.

  • best_epoch (int) – Epoch index at which the best validation loss was achieved.

  • best_val_loss (float)

fit(train_ds, val_ds, epochs=100)[source]#

Train the network.

Parameters:

  • train_ds (EMDataset) – Training split.

  • val_ds (EMDataset) – Validation split.

  • epochs (int) – Maximum number of epochs.

Return type:

self

get_weights()[source]#

Return model weights as numpy dict.

Return type:

dict[str, ndarray]

load_weights(weights)[source]#

Restore weights from numpy dict.

Parameters:

weights (dict[str, ndarray])

Return type:

None