xnn.common.finetune.heads.MultiHead#

class xnn.common.finetune.heads.MultiHead(model, heads)[source]#

Bases: InteratomicPotential

A potential with several readout heads over one shared trunk.

Head 0 is the wrapped model’s own readout (and reference energies); the other heads are copies of the modules the model names in its head_modules. Every head sees the same trunk, so a forward pass on a batch mixing heads runs the trunk once per head present (on that head’s structures) with the head’s parameters substituted, and nothing is computed twice.

Parameters:

  • model (InteratomicPotential) – The potential to give several heads. Its current readout becomes head 0.

  • heads (sequence of str) – The head names, in index order (["pt_head", "Default"] for the MACE replay convention: the pretrained readout serves the replay data, the copy adapts to the target data).

Variables:
Raises:

  • TypeError – If the model declares no head_modules.

  • ValueError – For an empty or repeated head list.

property species#

The trunk’s supported atomic numbers, if it lists them.

index(head)[source]#

The index of a head given by name or index.

Raises:

KeyError – For an unknown head name or an index out of range.

Parameters:

head (int | str)

Return type:

int

forward(data)[source]#

Evaluate every structure with the head it is labelled with.

Parameters:

data (AtomicGraph) – The batch; data.head holds the head index per structure (None means head 0 throughout).

Returns:

The trunk’s outputs (node_energy, energy, node_features, …) in the batch’s order.

Return type:

dict of str to torch.Tensor

using(head)[source]#

Make the trunk carry the modules of head for the duration of a block.

Lets any single-head routine (reference-energy estimation, pseudolabelling, a deployment export) act on one head through model; the trunk’s own head is restored afterwards.

Parameters:

head (int or str) – The head to swap in.

select(head)[source]#

A standalone single-head potential: the trunk with one head’s modules.

Parameters:

head (int or str) – The head to keep.

Returns:

A deep copy of the trunk carrying that head; the same class as the wrapped model, so every deployment path (TorchScript export, ASE calculator, MDI) applies unchanged.

Return type:

InteratomicPotential

set_atomic_energies(head, values)[source]#

Set the per-element reference energies of one head.

Parameters:

  • head (int or str) – The head.

  • values (mapping or sequence) – {Z: E0} (atomic numbers or symbols as keys) or one value per entry of the trunk’s species; see set_atomic_energies().

Return type:

None

atomic_energies(head)[source]#

The per-element reference energies of one head, {Z: E0}.

Parameters:

head (int | str)

Return type:

dict[int, float]

label(structures, head)[source]#

Tag structures with one of this model’s heads (see label_head()).

Parameters:

head (int | str)

set_use_fast(use_fast)[source]#

Forward the fast-path setting to the trunk (see xnn.common.models.fast.set_use_fast()).

Return type:

MultiHead

classmethod from_config(cfg)[source]#

Build through build_model() (extra["heads"]).

Return type:

MultiHead