train-deepmd.yml

# This is the configuration file for train stage with DeepMD-kit.
# The configuration contains 2 parts: executor context and workflow params.

executors:
  # This should correspond to the executor name in executor.yml
  hpc-cluster01:
    context:
      train:
        deepmd:
          # The template used to generate submit script.
          script_template:
            # Specify the header of the submit script,
            # Should be configured according to your system.
            header: |
              #SBATCH -N 1
              #SBATCH --ntasks-per-node=4
              #SBATCH --job-name=deepmd
              #SBATCH --partition=gpu

            # Specify the setup commands of the submit script,
            # Should be configured according to your system.
            setup: |
              set -e
              module load deepmd/2.2
              export OMP_NUM_THREADS=1
              set +e

          # Optional, specify the command to run deepmd-kit, default: dp.
          dp_cmd:

          # Optional, set this to true if you have multiple gpu cards in a batch job,  default: false
          multi_gpus_per_job:

workflow:
  train:
    deepmd:
      # The number of models to train.
      model_num: 4

      # Optional, whether to isolate the outliers in the training dataset.
      # Default: false
      isolate_outliers: true

      # Optional, the outlier cutoff for the training dataset.
      # Default: 10.
      outlier_f_cutoff: 5.

      # Optional, the weight of outlier data during training.
      # Default: 0.003
      outlier_weight: 0.003

      # Optional, the initial dataset to train the first model.
      # if not specified, it will use the dataset from previous label stage.
      # The value should correspond to the artifact name in artifact.yml
      init_dataset: [ h2o-64-train ]

      # Optional, skip training stage and use fixture models instead.
      # This is useful for debugging or skipping training.
      # Support globs, for example:
      # [ /path/to/*/frozen_model.pb ]
      fixture_models: []

      # Optional, define a custom input modifier function to modify the input template.
      # Useful when you want to modify the input template dynamically.
      # For example, change auto_prob_style according to train data.
      # You can define directly, or use !load_text path/to/fn.py to load from a file.
      input_modifier_fn: |
        def input_modifier_fn(input_data):
            input_data['model']['descriptor']['auto_prob_style'] = 'auto'
            return input_data

      # Optional, extra options that will be passed to dp train command.
      dp_train_opts: "--skip-neighbor-stat"

      # The input template for deepmd-kit.
      # You can put the content of the input template here,
      # or you can put the input template in a file, for example, deepmd-input.json,
      # and then you can load the input template from the file by
      #
      # input_template: !load_yaml deepmd-input.json
      #
      # Here we put the input template directly in the configuration file.
      # For more information about the input template, please refer to DeepMD-kit documentation.
      input_template:
        model:
          descriptor:
            type: se_a
            sel: [ 100, 100 ]
            rcut_smth: 0.5
            rcut: 5.0
            neuron: [ 25, 50, 100 ]
            resnet_dt: false
            axis_neuron: 16
            seed: 1
          fitting_net:
            neuron: [ 240, 240, 240]
            resnet_dt: true
            seed: 1
        learning_rate:
          type: exp
          start_lr: 0.001
          decay_steps: 2000
        loss:
          start_pref_e: 0.02
          limit_pref_e: 2
          start_pref_f: 1000
          limit_pref_f: 1
          start_pref_v: 0
          limit_pref_v: 0
        training:
          numb_steps: 10000
          seed: 1
          disp_file: lcurve.out
          disp_freq: 1000
          save_freq: 1000
          save_ckpt: model.ckpt
          disp_training: true
          time_training: true
          profiling: false
          profiling_file: timeline.json