pdft.training.train_basis_batched

pdft.training.train_basis_batched#

pdft.training.train_basis_batched(basis, *, dataset, loss, epochs, batch_size, optimizer='adam', validation_split=0.0, early_stopping_patience=5, warmup_frac=0.05, lr_peak=0.01, lr_final=0.001, max_grad_norm=None, shuffle=True, seed=0, val_every_k_epochs=1, frozen_indices=None)[source]#

Multi-image, multi-epoch trainer with cosine LR schedule.

Mirror of ParametricDFT.jl/src/training.jl::_train_basis_core (main). See parameter docs in the original training.py docstring.

Parameters:

  • frozen_indices (list[int] | None, optional) –

    List of integer indices into basis.tensors. Tensors at these indices are NOT updated during training — they stay at their initial values throughout. Each step computes gradients on all tensors normally; the update is then suppressed for frozen indices BEFORE any optimizer state is mutated (so Adam’s moment buffers for frozen indices remain zero). Useful for experiments that train only a subset of a circuit’s gates.

    Validation: all indices must satisfy 0 <= i < len(basis.tensors), no duplicates are allowed, and an empty list is treated as None (no-op). ValueError is raised on mis-specification.

  • dataset (Sequence)

  • loss (AbstractLoss)

  • epochs (int)

  • batch_size (int)

  • validation_split (float)

  • early_stopping_patience (int)

  • warmup_frac (float)

  • lr_peak (float)

  • lr_final (float)

  • max_grad_norm (float | None)

  • shuffle (bool)

  • seed (int)

  • val_every_k_epochs (int)

Return type:

TrainingResult