# lipengfei1409/deepmd-kit architecture

> Architecture diagram of the GitHub repository lipengfei1409/deepmd-kit, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/lipengfei1409/deepmd-kit
- Repository: https://github.com/lipengfei1409/deepmd-kit

## Overview

DeePMD-kit builds machine-learning interatomic potentials from reference data, supports training from scratch or fine-tuning pretrained checkpoints, then validates and exports models for simulation. Its main workflow spans data preparation, backend-specific training, model evaluation and optimization, and deployment through Python/native APIs or molecular-dynamics engines. Model architectures support energy and force prediction as well as properties such as dipoles, spin, and density of states. The graph reflects README-documented capabilities and sampled entrypoints; unsampled internals and integration wiring are kept general.

Main entry points: [`deepmd/main.py`](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/main.py) (CLI dispatch), [`deepmd/pt/train/training.py`](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/train/training.py) (Optimization), [`deepmd/infer/deep_eval.py`](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/infer/deep_eval.py) (Model inference), [`deepmd/pt/entrypoints/main.py`](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/entrypoints/main.py) (PyTorch training)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_workflow["Workflow entrypoints"]
  node_cli["CLI dispatch<br/>[main.py]"]
  node_pretrained["Pretrained models"]
  node_adapt["DPA-ADAPT<br/>[main.py]"]
end

subgraph group_training["Training and models"]
  node_data["Training data<br/>[dataset.py]"]
  node_pt_train["PyTorch training<br/>[main.py]"]
  node_jax_train["JAX training<br/>[main.py]"]
  node_pd_train["Paddle training<br/>[main.py]"]
  node_expt_train["Experimental training<br/>[main.py]"]
  node_architecture["Model architecture<br/>[dp_model.py]"]
  node_descriptor["Atomic descriptors"]
  node_tasks["Prediction heads"]
  node_loss["Training losses"]
  node_trainer["Optimization<br/>[training.py]"]
end

subgraph group_inference["Inference and validation"]
  node_eval["Model inference<br/>[deep_eval.py]"]
  node_test["Model testing<br/>[main.py]"]
  node_convert["Backend conversion<br/>[convert_backend.py]"]
  node_compress["Model compression<br/>[compress.py]"]
  node_freeze["Freeze and export<br/>[main.py]"]
end

subgraph group_deployment["Deployment integrations"]
  node_api["Python and native APIs<br/>[calculator.py]"]
  node_md["MD engines"]
  node_ase["ASE integration<br/>[ase_calc.py]"]
end

node_user(("Model user"))

node_user -->|"runs commands"| node_cli
node_cli -->|"resolves checkpoints"| node_pretrained
node_cli -->|"dispatches training"| node_pt_train
node_cli -->|"dispatches training"| node_jax_train
node_cli -->|"dispatches training"| node_pd_train
node_cli -->|"dispatches training"| node_expt_train
node_cli -->|"dispatches tests"| node_test
node_cli -->|"dispatches conversion"| node_convert
node_cli -->|"dispatches compression"| node_compress
node_cli -->|"dispatches export"| node_freeze
node_pretrained -->|"initializes fine-tuning"| node_pt_train
node_user -->|"starts adaptation"| node_adapt
node_adapt -->|"configures training"| node_pt_train
node_data -->|"supplies batches"| node_pt_train
node_pt_train -->|"runs optimization"| node_trainer
node_jax_train -->|"runs optimization"| node_trainer
node_pd_train -->|"runs optimization"| node_trainer
node_expt_train -->|"runs optimization"| node_trainer
node_trainer -->|"updates parameters"| node_architecture
node_architecture -->|"computes representations"| node_descriptor
node_architecture -->|"predicts targets"| node_tasks
node_tasks -->|"produces predictions"| node_loss
node_trainer -->|"optimizes against"| node_loss
node_user -->|"loads model"| node_eval
node_eval -->|"evaluates predictions"| node_test
node_eval -->|"serves predictions"| node_api
node_convert -->|"converts model"| node_eval
node_compress -->|"optimizes model"| node_eval
node_freeze -->|"exports model"| node_eval
node_api -->|"serves simulation"| node_md
node_api -->|"supports calculations"| node_ase

click node_cli "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/main.py"
click node_pretrained "https://github.com/lipengfei1409/deepmd-kit/tree/master/deepmd/pretrained"
click node_adapt "https://github.com/lipengfei1409/deepmd-kit/blob/master/dpa_adapt/main.py"
click node_data "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/utils/dataset.py"
click node_pt_train "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/entrypoints/main.py"
click node_jax_train "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/jax/entrypoints/main.py"
click node_pd_train "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pd/entrypoints/main.py"
click node_expt_train "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt_expt/entrypoints/main.py"
click node_architecture "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/model/model/dp_model.py"
click node_descriptor "https://github.com/lipengfei1409/deepmd-kit/tree/master/deepmd/pt/model/descriptor"
click node_tasks "https://github.com/lipengfei1409/deepmd-kit/tree/master/deepmd/pt/model/task"
click node_loss "https://github.com/lipengfei1409/deepmd-kit/tree/master/deepmd/pt/loss"
click node_trainer "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/train/training.py"
click node_eval "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/infer/deep_eval.py"
click node_test "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/entrypoints/main.py"
click node_convert "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/entrypoints/convert_backend.py"
click node_compress "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/entrypoints/compress.py"
click node_freeze "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/entrypoints/main.py"
click node_api "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/calculator.py"
click node_ase "https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/utils/ase_calc.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_cli,node_pretrained,node_adapt,node_user toneBlue
class node_data,node_pt_train,node_jax_train,node_pd_train,node_expt_train,node_architecture,node_descriptor,node_tasks,node_loss,node_trainer toneAmber
class node_eval,node_test,node_convert,node_compress,node_freeze toneMint
class node_api,node_md,node_ase toneRose
```

## Components

### Workflow entrypoints

- [CLI dispatch](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/main.py): `deepmd/main.py`
- [Pretrained models](https://github.com/lipengfei1409/deepmd-kit/tree/master/deepmd/pretrained): `deepmd/pretrained`
- [DPA-ADAPT](https://github.com/lipengfei1409/deepmd-kit/blob/master/dpa_adapt/main.py): `dpa_adapt/main.py`

### Training and models

- [Training data](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/utils/dataset.py): `deepmd/pt/utils/dataset.py`
- [PyTorch training](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/entrypoints/main.py): `deepmd/pt/entrypoints/main.py`
- [JAX training](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/jax/entrypoints/main.py): `deepmd/jax/entrypoints/main.py`
- [Paddle training](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pd/entrypoints/main.py): `deepmd/pd/entrypoints/main.py`
- [Experimental training](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt_expt/entrypoints/main.py): `deepmd/pt_expt/entrypoints/main.py`
- [Model architecture](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/model/model/dp_model.py): `deepmd/pt/model/model/dp_model.py`
- [Atomic descriptors](https://github.com/lipengfei1409/deepmd-kit/tree/master/deepmd/pt/model/descriptor): `deepmd/pt/model/descriptor`
- [Prediction heads](https://github.com/lipengfei1409/deepmd-kit/tree/master/deepmd/pt/model/task): `deepmd/pt/model/task`
- [Training losses](https://github.com/lipengfei1409/deepmd-kit/tree/master/deepmd/pt/loss): `deepmd/pt/loss`
- [Optimization](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/train/training.py): `deepmd/pt/train/training.py`

### Inference and validation

- [Model inference](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/infer/deep_eval.py): `deepmd/infer/deep_eval.py`
- [Model testing](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/entrypoints/main.py): `deepmd/entrypoints/main.py`
- [Backend conversion](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/entrypoints/convert_backend.py): `deepmd/entrypoints/convert_backend.py`
- [Model compression](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/entrypoints/compress.py): `deepmd/pt/entrypoints/compress.py`
- [Freeze and export](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/entrypoints/main.py): `deepmd/pt/entrypoints/main.py`

### Deployment integrations

- [Python and native APIs](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/calculator.py): `deepmd/calculator.py`
- **MD engines**
- [ASE integration](https://github.com/lipengfei1409/deepmd-kit/blob/master/deepmd/pt/utils/ase_calc.py): `deepmd/pt/utils/ase_calc.py`

### Other components

- **Model user**

## Connections

- Model user → CLI dispatch: runs commands
- CLI dispatch → Pretrained models: resolves checkpoints
- CLI dispatch → PyTorch training: dispatches training
- CLI dispatch → JAX training: dispatches training
- CLI dispatch → Paddle training: dispatches training
- CLI dispatch → Experimental training: dispatches training
- CLI dispatch → Model testing: dispatches tests
- CLI dispatch → Backend conversion: dispatches conversion
- CLI dispatch → Model compression: dispatches compression
- CLI dispatch → Freeze and export: dispatches export
- Pretrained models → PyTorch training: initializes fine-tuning
- Model user → DPA-ADAPT: starts adaptation
- DPA-ADAPT → PyTorch training: configures training
- Training data → PyTorch training: supplies batches
- PyTorch training → Optimization: runs optimization
- JAX training → Optimization: runs optimization
- Paddle training → Optimization: runs optimization
- Experimental training → Optimization: runs optimization
- Optimization → Model architecture: updates parameters
- Model architecture → Atomic descriptors: computes representations
- Model architecture → Prediction heads: predicts targets
- Prediction heads → Training losses: produces predictions
- Optimization → Training losses: optimizes against
- Model user → Model inference: loads model
- Model inference → Model testing: evaluates predictions
- Model inference → Python and native APIs: serves predictions
- Backend conversion → Model inference: converts model
- Model compression → Model inference: optimizes model
- Freeze and export → Model inference: exports model
- Python and native APIs → MD engines: serves simulation
- Python and native APIs → ASE integration: supports calculations