deepmd.pt_expt.entrypoints.main

deepmd.pt_expt.entrypoints.main#

Training entrypoint for the pt_expt backend.

Attributes#

`log`
`_PT_EXPT_MODEL_SUFFIXES`

Classes#

`SummaryPrinter`	Summary printer for pt_expt.
`PTExptTrainEntrypoint`	pt_expt implementation of the common training entrypoint pipeline.

Functions#

`_ensure_pt_expt_model_suffix`(→ str \| None)	Append the default checkpoint suffix when a model path is a prefix.
`_update_changed_model_tensors`(→ None)	Copy changed tensors into an existing state dict without breaking aliases.
`_detect_lmdb_path`(→ str \| None)	Return the LMDB path when `systems_raw` is a scalar LMDB string.
`_get_neighbor_stat_data`(→ Any)	Return a data proxy suitable for `BaseModel.update_sel` (neighbor stat).
`_build_data_system`(...)	Build a data system from dataset config, routing LMDB paths to LmdbDataSystem.
`get_trainer`(→ deepmd.pt_expt.train.training.Trainer)	Build a `training.Trainer` from a normalised config.
`train`(→ None)	Run training with the pt_expt backend.
`_default_output_path`(→ str)	Default a suffixless frozen-model path from the RESOLVED lower kind.
`freeze`(→ None)	Freeze a pt_expt checkpoint into a .pte/.pt2 exported model.
`change_bias`(→ None)	Change the output bias of a pt_expt model.
`main`(→ None)	Entry point for the pt_expt backend CLI.

Module Contents#

deepmd.pt_expt.entrypoints.main.log[source]#

deepmd.pt_expt.entrypoints.main._PT_EXPT_MODEL_SUFFIXES = ('.pt', '.pte', '.pt2')[source]#

deepmd.pt_expt.entrypoints.main._ensure_pt_expt_model_suffix(model_path: str | None) → str | None[source]#: Append the default checkpoint suffix when a model path is a prefix.

deepmd.pt_expt.entrypoints.main._update_changed_model_tensors(target_state_dict: dict[str, Any], source_state_dict: dict[str, Any]) → None[source]#: Copy changed tensors into an existing state dict without breaking aliases.

deepmd.pt_expt.entrypoints.main._detect_lmdb_path(systems_raw: Any) → str | None[source]#

Return the LMDB path when systems_raw is a scalar LMDB string.

Returns None for non-LMDB inputs. Raises ValueError if systems_raw is a list containing any LMDB path, so both _get_neighbor_stat_data and _build_data_system fail with the same clear message instead of the opaque error from process_systems() / DeepmdData.

deepmd.pt_expt.entrypoints.main._get_neighbor_stat_data(dataset_params: dict[str, Any], type_map: list[str] | None) → Any[source]#

Return a data proxy suitable for BaseModel.update_sel (neighbor stat).

Routes a scalar LMDB systems path through dpmodel’s make_neighbor_stat_data; falls back to the legacy get_data for npy/HDF5 directories.

deepmd.pt_expt.entrypoints.main._build_data_system(dataset_params: dict[str, Any], type_map: list[str], seed: int | None = None, rank: int = 0, world_size: int = 1) → deepmd.utils.data_system.DeepmdDataSystem | deepmd.pt_expt.utils.lmdb_dataset.LmdbDataSystem[source]#

Build a data system from dataset config, routing LMDB paths to LmdbDataSystem.

A scalar systems value pointing at an LMDB directory triggers the LMDB adapter; otherwise we fall through to the legacy DeepmdDataSystem path with system expansion. rank and world_size shard LMDB training batches without changing legacy data systems.

deepmd.pt_expt.entrypoints.main.get_trainer(config: dict[str, Any], init_model: str | None = None, restart_model: str | None = None, finetune_model: str | None = None, finetune_links: dict | None = None, shared_links: dict | None = None) → deepmd.pt_expt.train.training.Trainer[source]#: Build a training.Trainer from a normalised config.

class deepmd.pt_expt.entrypoints.main.SummaryPrinter[source]#

Bases: deepmd.utils.summary.SummaryPrinter

Summary printer for pt_expt.

is_built_with_cuda() → bool[source]#: Check if PyTorch was built with CUDA.

is_built_with_rocm() → bool[source]#: Check if PyTorch was built with ROCm.

get_compute_device() → str[source]#: Get the selected compute device.

get_ngpus() → int[source]#: Get the number of visible CUDA devices.

get_backend_info() → dict[source]#: Get backend information.

get_device_name() → str | None[source]#: Return the current CUDA device name when available.

class deepmd.pt_expt.entrypoints.main.PTExptTrainEntrypoint[source]#

Bases: deepmd.dpmodel.train.AbstractTrainEntrypoint

pt_expt implementation of the common training entrypoint pipeline.

finetune_links: dict[str, Any] | None = None[source]#

shared_links: dict[str, Any] | None = None[source]#

_owns_process_group = False[source]#

prepare_options(options: deepmd.dpmodel.train.TrainEntrypointOptions) → deepmd.dpmodel.train.TrainEntrypointOptions[source]#: Normalize checkpoint prefixes accepted by the train CLI.

preprocess_config(config: dict[str, Any], options: deepmd.dpmodel.train.TrainEntrypointOptions) → dict[str, Any][source]#: Apply pt_expt multi-task, finetune, and pretrained-model preprocessing.

update_neighbor_stat(config: dict[str, Any], options: deepmd.dpmodel.train.TrainEntrypointOptions, *, multi_task: bool) → tuple[dict[str, Any], float | dict[str, float | None] | None][source]#: Update pt_expt descriptor selections from neighbor statistics.

print_summary() → None[source]#: Print pt_expt backend summary.

setup_run(options: deepmd.dpmodel.train.TrainEntrypointOptions, config: dict[str, Any]) → None[source]#: Initialize pt_expt distributed training when launched by torchrun/srun.

teardown_run(options: deepmd.dpmodel.train.TrainEntrypointOptions, config: dict[str, Any]) → None[source]#: Destroy the pt_expt distributed process group if this entrypoint made one.

run_training(config: dict[str, Any], options: deepmd.dpmodel.train.TrainEntrypointOptions, neighbor_stat: Any) → None[source]#: Build and run the pt_expt trainer.

deepmd.pt_expt.entrypoints.main.train(input_file: str, init_model: str | None = None, restart: str | None = None, finetune: str | None = None, model_branch: str = '', use_pretrain_script: bool = False, skip_neighbor_stat: bool = False, output: str = 'out.json') → None[source]#

Run training with the pt_expt backend.

Parameters:

input_filestr: Path to the JSON configuration file.
init_modelstr or None: Path to a checkpoint to initialise weights from.
restartstr or None: Path to a checkpoint to restart training from.
finetunestr or None: Path to a pretrained checkpoint to fine-tune from.
model_branchstr: Branch to select from a multi-task pretrained model.
use_pretrain_scriptbool: If True, copy descriptor/fitting params from the pretrained model.
skip_neighbor_statbool: Skip neighbour statistics calculation.
outputstr: Where to dump the normalised config.

deepmd.pt_expt.entrypoints.main._default_output_path(output: str, lower_kind: str) → str[source]#

Default a suffixless frozen-model path from the RESOLVED lower kind.

Parameters:

outputstr: Requested output path. An explicit .pte / .pt2 suffix is preserved; any other (or missing) suffix is replaced.
lower_kindstr: The resolved lower kind (after freeze’s native-spin resolution): "graph" selects .pt2 (an AOTI .pt2 archive is what the C++ graph path consumes), anything else .pte.

Returns:

str: The output path with a definite .pte / .pt2 suffix.

deepmd.pt_expt.entrypoints.main.freeze(model: str, output: str = 'frozen_model', head: str | None = None, lower_kind: str = 'nlist') → None[source]#

Freeze a pt_expt checkpoint into a .pte/.pt2 exported model.

Parameters:

modelstr: Path to the checkpoint file (.pt).
outputstr: Path for the output file. When it carries no .pte/.pt2 suffix, the suffix is chosen from the RESOLVED lower kind (after the native-spin resolution below): .pt2 for a graph lower, .pte otherwise.
headstr or None: Head to freeze in multi-task mode.
lower_kindstr: Lower-level export form: "nlist" (default, dense neighbor-list lower) or "graph" (NeighborGraph edge-list lower). "graph" is only valid for graph-eligible models (mixed_types and uses_graph_lower: dpa1/se_atten with concat type embedding) and selects the C++ graph inference path; the per-atom virial is enabled for it (near-free in the graph path: one extra scatter off the shared single backward). NOTE: for smooth_type_embedding=True the carry-all graph attention intentionally drops the dense layout’s sel-padding terms from the softmax denominator, so graph-form results are sel-independent and differ from the legacy dense lower by up to ~1e-4 (see DescrptDPA1.call_graph). A NATIVE-spin model (NativeSpinEnergyModel) implements ONLY the graph lower, so lower_kind is resolved to "graph" for it regardless of the requested value – before the export ABI and the default output suffix are selected.

deepmd.pt_expt.entrypoints.main.change_bias(input_file: str, mode: str = 'change', bias_value: list | None = None, datafile: str | None = None, system: str = '.', numb_batch: int = 0, model_branch: str | None = None, output: str | None = None) → None[source]#

Change the output bias of a pt_expt model.

Parameters:

input_filestr: Path to the model file (.pt checkpoint or .pte frozen model).
modestr: "change" or "set".
bias_valuelist or None: User-defined bias values (one per type).
datafilestr or None: File listing data system paths.
systemstr: Data system path (used when datafile is None).
numb_batchint: Number of batches for statistics (0 = all).
model_branchstr or None: Branch name for multi-task models.
outputstr or None: Output file path.

deepmd.pt_expt.entrypoints.main.main(args: list[str] | argparse.Namespace | None = None) → None[source]#

Entry point for the pt_expt backend CLI.

Parameters:

argslist[str] | argparse.Namespace | None: Command-line arguments or pre-parsed namespace.