TACE-Scripts Tutorial#

This section introduces some commonly used scripts for training, inference, and various utility tasks. The supported scripts are:

  • tace-train

  • tace-eval

  • tace-export

  • tace-finetune

  • tace-convert

  • tace-average

  • tace-update

  • tace-clean

  • tace-graph

  • tace-download

  • tace-split

Note

All scripts support the -h option to display all available command-line arguments and usage instructions.

  1. tace-train This script is used to start model training. Users can configure training parameters and initiate the training process.

tace-train -cn *.yaml

# In particular, if your YAML file is named 'tace.yaml', you can directly use:
tace-train

  1. tace-eval After training is complete, this script evaluates the model and produces an inference result file.

# You should at least include -i and -m options
# -i specifies the input file (ase.io.read) for inference
# -m specifies the model checkpoint file (.ckpt, .pt and .pth)
# If you want to test, add --test,
# don't forget to specify your key, if your property key is not the default value
# For other options, use -h

tace-eval -i *.xyz -m *.ckpt

  1. tace-export Conversion of model formats

    • ckpt:

      The .ckpt file is a checkpoint saved by the Lightning callback. It contains not only the model but also all related metadata, allowing the training process to be resumed.

    • pt/pth:

      Contains only the model itself without additional data. It is used for inference.

# Currently, only the use of LAMMPS requires mandatory model export.
# For other backends, you may also export the model if you wish to reduce its storage size.
tace-export -m *.ckpt --backend lammps

  1. tace-clean

# This script is used to clean the output files generated by TACE
tace-clean

  1. tace-split

# This script is used to split the dataset
# -n specify the number of configurations in each set
tace-split -i *.xyz -n 800 100 100 --seed 42

6. tace-graph This script is mainly used for pre-building graphs for large datasets (>1 million). It is particularly useful if you might repeat training. Otherwise, each time you retrain, you would need to rebuild the graph. The input file must be exactly the same as the one used for training. When using this script, it will stop immediately after building the graph. In practical use with tace-train, generally, if the dataset storage mode is lmdb and the specified shard_dirs contain pre-built graphs, it will automatically read them and skip the graph-building phase.

tace-graph -cn *.yaml

  1. tace-convert

This script is mainly used to perform architecture-level or parameter-level transformations on an existing model.

Typical use cases include (but are not limited to):

  • Merging LoRA parameters back into the base model

  1. tace-update

This script is used to update internal statistical information stored in an existing model.

Typical examples include:

  • Updating reference atomic energies

  • Adjusting internal scale and shift parameters

  1. tace-average

This script accepts a series of models with identical architectures and computes the average of all model parameters.

In practice, this serves as a manual implementation of Stochastic Weight Averaging (SWA).

A common workflow is:

  • Reduce the weight of force loss (i.e., increase the relative importance of energy) in the late stage of training

  • Continue training with a low learning rate

  • Collect a sequence of checkpoints

  • Apply tace-average to obtain an averaged model

This averaged model often exhibits better generalization performance.

Note

If average models is used, it is generally not recommended to use EMA at the same time, as both methods perform temporal averaging of parameters and may interfere with each other.

  1. tace-download

This script is used to download models or datasets provided by TACE.

However, manual downloading is generally recommended, as network connectivity to the server may be unreliable in some environments.

  1. tace-finetune

This script is used to generate configuration files for fine-tuning an existing model.

The generated finetune_config.yaml file can be further customized before launching the fine-tuning process().