File size: 7,180 Bytes
8e4e24e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
core:
  version: ${get_flowmm_version:}
  tags:
  - ${now:%Y-%m-%d}
  - ⚡️pytorch lightning
logging:
  val_check_interval: 25
  wandb:
    project: uniug_uni
    entity: null
    log_model: true
    mode: cloud
    experiment_name: mp20_PathRefine_wuns0.001_wc0.01_e400_e450
  wandb_watch:
    log: all
    log_freq: 500
  lr_monitor:
    logging_interval: step
    log_momentum: false
optim:
  optimizer:
    _target_: torch.optim.AdamW
    lr: 0.0005
    weight_decay: 0.0
  lr_scheduler:
    _target_: torch.optim.lr_scheduler.CosineAnnealingLR
    T_max: ${data.train_max_epochs}
    eta_min: 1.0e-05
  interval: epoch
  ema_decay: 0.999
train:
  deterministic: warn
  random_seed: 42
  pl_trainer:
    fast_dev_run: false
    strategy: ddp_find_unused_parameters_true
    num_nodes: 1
    devices: 8
    accelerator: gpu
    precision: 32
    max_epochs: ${data.train_max_epochs}
    accumulate_grad_batches: 1
    num_sanity_val_steps: 1
    gradient_clip_val: 10.0
    gradient_clip_algorithm: norm
    profiler: simple
    log_every_n_steps: 500
  monitor_metric: val/loss_csp
  monitor_metric_mode: min
  model_checkpoints:
    save_top_k: 1
    verbose: false
    save_last: false
  every_n_epochs_checkpoint:
    every_n_epochs: 25
    save_top_k: -1
    verbose: false
    save_last: false
val:
  compute_nll: false
test:
  compute_nll: false
  compute_loss: true
integrate:
  div_mode: rademacher
  method: euler
  num_steps: 1000
  normalize_loglik: true
  inference_anneal_slope: 0.0
  inference_anneal_offset: 0.0
base_distribution_from_data: false
partial_ckpt_load: true
partial_ckpt_path: /mnt/ai4sci_develop_fast/songyouli/crystal-uniug/runs/trash/uniug_uni/bz256_mp20_mptsub_tclearn_freeze_epoch500/every_n_epochs/epoch=399-step=754800.ckpt
data:
  dataset_name: mp_20
  dim_coords: 3
  root_path: ${oc.env:DATA_DIR}/mp_20
  prop: formation_energy_per_atom
  num_targets: 1
  niggli: true
  primitive: false
  graph_method: crystalnn
  lattice_scale_method: scale_length
  preprocess_workers: 30
  readout: mean
  max_atoms: 20
  otf_graph: false
  eval_model_name: mp20
  tolerance: 0.1
  use_space_group: false
  use_pos_index: false
  train_max_epochs: 500
  early_stopping_patience: 100000
  teacher_forcing_max_epoch: 500
  md_dataset_name: mptsubmp20_v0
  root_path_md: ${oc.env:DATA_DIR}/${data.md_dataset_name}
  require_order: false
  prop_md:
  - energy
  - forces
  energy_only: true
  t_mode: c_learn
  t_constant: 1.0
  datamodule:
    _target_: uniug.datamodule_uni.CrystDataModule
    task_mode: uni
    datasets:
      train:
        _target_: diffcsp.pl_data.dataset.CrystDataset
        name: Formation energy train
        path: ${data.root_path}/train.csv
        save_path: ${data.root_path}/train_ori.pt
        prop: ${data.prop}
        niggli: ${data.niggli}
        primitive: ${data.primitive}
        graph_method: ${data.graph_method}
        tolerance: ${data.tolerance}
        use_space_group: ${data.use_space_group}
        use_pos_index: ${data.use_pos_index}
        lattice_scale_method: ${data.lattice_scale_method}
        preprocess_workers: ${data.preprocess_workers}
      val:
      - _target_: diffcsp.pl_data.dataset.CrystDataset
        name: Formation energy val
        path: ${data.root_path}/val.csv
        save_path: ${data.root_path}/val_ori.pt
        prop: ${data.prop}
        niggli: ${data.niggli}
        primitive: ${data.primitive}
        graph_method: ${data.graph_method}
        tolerance: ${data.tolerance}
        use_space_group: ${data.use_space_group}
        use_pos_index: ${data.use_pos_index}
        lattice_scale_method: ${data.lattice_scale_method}
        preprocess_workers: ${data.preprocess_workers}
      test:
      - _target_: diffcsp.pl_data.dataset.CrystDataset
        name: Formation energy test
        path: ${data.root_path}/test.csv
        save_path: ${data.root_path}/test_ori.pt
        prop: ${data.prop}
        niggli: ${data.niggli}
        primitive: ${data.primitive}
        graph_method: ${data.graph_method}
        tolerance: ${data.tolerance}
        use_space_group: ${data.use_space_group}
        use_pos_index: ${data.use_pos_index}
        lattice_scale_method: ${data.lattice_scale_method}
        preprocess_workers: ${data.preprocess_workers}
      train_md:
        _target_: diffcsp.pl_data.dataset.CrystDataset
        name: Formation energy train
        path: ${data.root_path_md}/train.csv
        save_path: ${data.root_path_md}/train_ori.pt
        require_order: ${data.require_order}
        prop: ${data.prop_md}
        niggli: ${data.niggli}
        primitive: ${data.primitive}
        graph_method: ${data.graph_method}
        tolerance: ${data.tolerance}
        use_space_group: ${data.use_space_group}
        use_pos_index: ${data.use_pos_index}
        lattice_scale_method: ${data.lattice_scale_method}
        preprocess_workers: ${data.preprocess_workers}
      val_md:
      - _target_: diffcsp.pl_data.dataset.CrystDataset
        name: Formation energy val
        path: ${data.root_path_md}/val.csv
        save_path: ${data.root_path_md}/val_ori.pt
        require_order: ${data.require_order}
        prop: ${data.prop_md}
        niggli: ${data.niggli}
        primitive: ${data.primitive}
        graph_method: ${data.graph_method}
        tolerance: ${data.tolerance}
        use_space_group: ${data.use_space_group}
        use_pos_index: ${data.use_pos_index}
        lattice_scale_method: ${data.lattice_scale_method}
        preprocess_workers: ${data.preprocess_workers}
    num_workers:
      train: 40
      val: 40
      test: 40
    batch_size:
      train: 32
      val: 32
      test: 32
      train_md: 32
      val_md: 32
    pin_memory: true
    persistent_workers: false
    prefetch_factor: 2
model:
  w_csp: 0.95
  w_pflow: 0.2
  w_time: 0.1
  cost_coord: 400.0
  cost_lattice: 1.0
  cost_type: 0.0
  cost_energy: 1.0
  cost_forces: 1.0
  cost_stress: 1.0
  affine_combine_costs: true
  target_distribution: conditional
  self_cond: false
  t_pflow_clip: false
  use_tangent: false
  tclearn_freeze_epoch: 250
  use_uns_flow_task: true
  w_uns_flow: 0.001
  w_consist: 0.01
  uns_flow_tscale: true
  uns_flow_tlearn: true
  uns_path_refine: true
  uns_path_t_clip: 0.9
  use_consist_flow: true
  consist_freeze_epoch: 450
  uns_flow_freeze_epoch: 400
  manifold_getter:
    atom_type_manifold: null_manifold
    coord_manifold: flat_torus_01
    lattice_manifold: lattice_params
    length_inner_coef: 1.0
vectorfield:
  _target_: uniug.arch_uni.FlowmmUniModel
  force_pred_way: direct
  use_pflow_head: true
  hidden_dim: 512
  time_dim: 256
  num_layers: 6
  act_fn: silu
  dis_emb: sin
  num_freqs: 128
  edge_style: fc
  max_neighbors: 20
  cutoff: 7.0
  ln: true
  use_log_map: true
  dim_atomic_rep: ${get_dim_atomic_rep:${model.manifold_getter.atom_type_manifold}}
  lattice_manifold: ${model.manifold_getter.lattice_manifold}
  concat_sum_pool: true
  represent_num_atoms: true
  represent_angle_edge_to_lattice: true
  self_edges: false
  self_cond: ${model.self_cond}
  t_mode: ${data.t_mode}
  t_mask: -1.0
  use_pflow_task: false
  tlearn_clip: false
  learnable_time_emb: false