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,613 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
142
143
144
145
146
147
# yamllint disable rule:line-length
# Smoke-test config for NequIP on OneScience.
# This config uses a tiny synthetic Cu dataset and a small model to verify the
# training pipeline (data loading, model build, forward, backward, checkpoint).

run: [train, test]

cutoff_radius: 4.0

num_layers: 2
l_max: 1
num_features: 8

model_type_names: [Cu]
chemical_species: ${model_type_names}

monitored_metric: val0_epoch/weighted_sum

# ============
#     DATA
# ============
data:
  _target_: onescience.datapipes.materials.nequip.datamodule.ASEDataModule
  seed: 456

  split_dataset:
    file_path: ${demo_dir:reference_data/smoke.xyz}
    train: 0.75
    val: 0.125
    test: 0.125

  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: 2
    num_workers: 0
    shuffle: true
  val_dataloader:
    _target_: torch.utils.data.DataLoader
    batch_size: 2
    num_workers: 0
  test_dataloader: ${data.val_dataloader}

  stats_manager:
    _target_: onescience.datapipes.materials.nequip.CommonDataStatisticsManager
    dataloader_kwargs:
      batch_size: 2
    type_names: ${model_type_names}

# =============
#    TRAINER
# =============
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: true

  callbacks:
    - _target_: lightning.pytorch.callbacks.ModelCheckpoint
      monitor: ${monitored_metric}
      dirpath: ${hydra:runtime.output_dir}/checkpoints
      filename: best
      save_last: true

# =====================
#    TRAINING MODULE
# =====================
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
    seed: 456
    model_dtype: float32
    type_names: ${model_type_names}
    r_max: ${cutoff_radius}
    num_bessels: 4
    bessel_trainable: false
    polynomial_cutoff_p: 6
    num_layers: ${num_layers}
    l_max: ${l_max}
    parity: false
    num_features: ${num_features}
    radial_mlp_depth: 1
    radial_mlp_width: 16
    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

# Slurm / launch metadata used by demo/run.sh
name: nequip_smoke
launch:
  mode: local
  num_nodes: 1
  num_gpus: 1
slurm:
  partition: hx1hdnormal01
  nodelist: a01r1n02
  time: "00:10:00"
  cpus_per_task: 8
env:
  OMP_NUM_THREADS: 1