| |
| |
| |
| 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 |
|
|
