Hugging Face's logo

OneScience-Group
/
NequIP

English
OneScience
NequIP
machine-learning-potential
molecular-simulation
materials-computing
graph-neural-network
equivariant-neural-network
training
fine-tuning
inference
Model card Files
xet
 Community
NequIP
File size: 4,111 Bytes
3e02ab8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
# Eight-DCU DDP training plan for the official NequIP fcu.xyz tutorial.
# DataLoader batch size is per rank: batch_size=5 gives a global batch of 40.
# The learning rate remains conservative at the official 0.01; tune it for
# production based on validation behavior rather than scaling it blindly.
run: [train, test]

cutoff_radius: 5.0
num_layers: 4
l_max: 1
num_features: 32
model_type_names: [C, H, O, Cu]
chemical_species: ${model_type_names}
monitored_metric: val0_epoch/weighted_sum

data:
  _target_: onescience.datapipes.materials.nequip.datamodule.ASEDataModule
  seed: 456
  split_dataset:
    file_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/matchem/NequIP/fcu.xyz
    train: 0.8
    val: 0.1
    test: 0.1
  transforms:
    - _target_: onescience.datapipes.materials.nequip.transforms.ChemicalSpeciesToAtomTypeMapper
      model_type_names: ${model_type_names}
    - _target_: onescience.datapipes.materials.nequip.transforms.NeighborListTransform
      r_max: ${cutoff_radius}
  train_dataloader:
    _target_: torch.utils.data.DataLoader
    batch_size: 5
    num_workers: 0
    shuffle: true
  val_dataloader:
    _target_: torch.utils.data.DataLoader
    batch_size: 10
    num_workers: 0
  test_dataloader: ${data.val_dataloader}
  stats_manager:
    _target_: onescience.datapipes.materials.nequip.CommonDataStatisticsManager
    dataloader_kwargs:
      batch_size: 10
    type_names: ${model_type_names}

trainer:
  _target_: lightning.Trainer
  accelerator: gpu
  devices: 8
  num_nodes: 1
  strategy:
    _target_: lightning.pytorch.strategies.DDPStrategy
  enable_checkpointing: true
  max_epochs: 1000
  max_time: "03:00:00:00"
  log_every_n_steps: 20
  logger:
    _target_: lightning.pytorch.loggers.CSVLogger
    save_dir: ${hydra:runtime.output_dir}
    name: metrics
    version: 0
  enable_progress_bar: true
  callbacks:
    - _target_: onescience.utils.nequip.train.callbacks.PlainTextMetricsLogger
    - _target_: lightning.pytorch.callbacks.EarlyStopping
      monitor: ${monitored_metric}
      min_delta: 1e-3
      patience: 20
    - _target_: lightning.pytorch.callbacks.ModelCheckpoint
      monitor: ${monitored_metric}
      dirpath: ${hydra:runtime.output_dir}/checkpoints
      filename: best
      save_last: true

training_module:
  _target_: onescience.utils.nequip.train.EMALightningModule
  ema_decay: 0.999
  loss:
    _target_: onescience.utils.nequip.train.EnergyForceLoss
    per_atom_energy: true
    coeffs:
      total_energy: 1.0
      forces: 1.0
  val_metrics:
    _target_: onescience.utils.nequip.train.EnergyForceMetrics
    coeffs:
      total_energy_mae: 1.0
      forces_mae: 1.0
  train_metrics: ${training_module.val_metrics}
  test_metrics: ${training_module.val_metrics}
  optimizer:
    _target_: torch.optim.Adam
    lr: 0.01
  lr_scheduler:
    scheduler:
      _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
      factor: 0.6
      patience: 5
      threshold: 0.2
      min_lr: 1.0e-6
    monitor: ${monitored_metric}
    interval: epoch
    frequency: 1
  model:
    _target_: onescience.models.nequip.model.NequIPGNNModel
    compile_mode: eager
    seed: 456
    model_dtype: float32
    type_names: ${model_type_names}
    r_max: ${cutoff_radius}
    num_bessels: 8
    bessel_trainable: false
    polynomial_cutoff_p: 6
    num_layers: ${num_layers}
    l_max: ${l_max}
    parity: true
    num_features: ${num_features}
    radial_mlp_depth: 2
    radial_mlp_width: 64
    avg_num_neighbors: ${training_data_stats:num_neighbors_mean}
    per_type_energy_scales: ${training_data_stats:per_type_forces_rms}
    per_type_energy_shifts: ${training_data_stats:per_atom_energy_mean}
    per_type_energy_scales_trainable: false
    per_type_energy_shifts_trainable: false
    pair_potential:
      _target_: onescience.models.nequip.nn.pair_potential.ZBL
      units: metal
      chemical_species: ${chemical_species}

name: nequip_fcu_tutorial_8dcu
launch:
  mode: auto
  num_nodes: 1
  num_gpus: 8
slurm:
  partition: hx1hdnormal01
  nodelist: ""
  time: "3-00:00:00"
  cpus_per_task: 8
env:
  OMP_NUM_THREADS: 8
  NCCL_DEBUG: "WARN"