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: 3,706 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
# NequIP 0.19 tutorial reproduction using the official fcu.xyz dataset.
# The validated smoke schedule uses two complete epochs. Set max_epochs to 1000
# to match the upstream tutorial's full training schedule.
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: 1
  num_nodes: 1
  enable_checkpointing: true
  max_epochs: 2
  log_every_n_steps: 1
  logger: false
  enable_progress_bar: false
  callbacks:
    - _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: 1e-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
launch:
  mode: local
  num_nodes: 1
  num_gpus: 1
slurm:
  partition: hx1hdnormal01
  nodelist: a01r1n02
  time: "00:30:00"
  cpus_per_task: 8
env:
  OMP_NUM_THREADS: 8