6.1. Freeze a model#
The trained neural network is extracted from a checkpoint and dumped into a model file. This process is called “freezing” a model. To freeze a model, typically one does
dp freeze -o model.pb
in the folder where the model is trained. The output model is called model.pb. The idea and part of our code are from Morgan.
dp --tf2 freeze -c model.ckpt -o model.savedmodeltf
When -c names a checkpoint prefix, the backend also checks the corresponding path with .tf2 appended, so the example reads model.ckpt.tf2 and writes the TensorFlow SavedModel to model.savedmodeltf. If -c is omitted, it defaults to the current directory. For a multi-task checkpoint, select a branch with --head CHOSEN_BRANCH. SavedModel export requires graph-traceable model code; descriptor-specific export requirements are documented on the corresponding model pages.
dp --pt freeze -o model.pth
in the folder where the model is trained. The output model is called model.pth.
In multi-task mode, you need to choose one available heads (e.g. CHOSEN_BRANCH) by --head to specify which model branch you want to freeze:
dp --pt freeze -o model_branch1.pth --head CHOSEN_BRANCH
The output model is called model_branch1.pth, which is the specifically frozen model with the CHOSEN_BRANCH head.
dp --pt-expt freeze -c model.ckpt.pt -o model
The backend writes .pte for the dense neighbor-list lower form and .pt2 for the graph lower form. A suffixless output lets DeePMD-kit select the matching extension. --lower-kind graph requires a graph-eligible model. Conversely, a graph-capable DPA model may override a requested nlist lower with the graph form and emit a warning. In multi-task mode, select a model branch with --head CHOSEN_BRANCH.
dp --pd freeze -o model
in the folder where the model is trained. The output model is called model.json and model.pdiparams.
In multi-task mode, you need to choose one available heads (e.g. CHOSEN_BRANCH) by --head to specify which model branch you want to freeze:
dp --pd freeze -o model_branch1 --head CHOSEN_BRANCH
The output model is called model_branch1.json, which is the specifically frozen model with the CHOSEN_BRANCH head.
dp --jax freeze -c model.ckpt.jax -o model.hlo
The JAX backend can write a StableHLO .hlo model, a lossless .jax serialization, or a JAX2TF .savedmodel model. The .hlo and .savedmodel formats work with the normal dp test/DeepPot route; .jax is intended for checkpoint round-tripping and JAX-MD and is not a DeepEval model format. The .savedmodel format requires TensorFlow and is the JAX format that supports the C++ inference interface.
6.1.1. Freeze a JAX model with Hessian output
#
Use --hessian to add coordinate-Hessian output to a frozen JAX energy model:
dp --jax freeze -c model.ckpt.jax -o model-hessian.hlo --hessian
The option applies to JAX .hlo, .jax, and .savedmodel outputs. A model whose serialized definition already enables Hessian mode retains that mode even when --hessian is omitted.