xnn.common.finetune.replay.pseudolabel#

xnn.common.finetune.replay.pseudolabel(model, structures, batch_size=32, device=None, stress=None, head=None)[source]#

Label structures with a model’s own energies, forces (and stress).

Parameters:

model (torch.nn.Module) – The pretrained potential: bare, wrapped in ForceStressOutput, or a MultiHead (then head picks the head that labels; by default the first).
structures (iterable of dict) – The structures to relabel; other keys (head, weight, cell) are kept.
batch_size (int, optional) – Structures per forward pass, by default 32.
device (torch.device, optional) – Where to evaluate; by default where the model is.
stress (bool, optional) – Whether to label the stress; by default when any structure has a cell.
head (int or str, optional) – The head of a multi-head model that produces the labels.

Returns:

Copies of the structures with energy, forces (and stress) replaced by the predictions.

Return type:

list of dict