File size: 1,585 Bytes
3a8e4d3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
# @package _global_

# All parameters below will be merged with parameters from default configurations set above 
# this allows you to overwrite only specified parameters

# Options for training:
# - /data: [pdb, pdb_sp, pdb_sp_afesm, pdb_sp_afesme]
# - /model/architecture: 
#         [foldingdit_100M, foldingdit_360M, foldingdit_700M, 
#          foldingdit_1.1B, foldingdit_1.6B, foldingdit_3B]
# - /trainer: [default, fsdp]

defaults:
  - override /data: pdb
  - override /model: simplefold
  - override /model/architecture: foldingdit_100M
  - override /model/sampler: euler_maruyama
  - override /model/processor: protein_processor
  - override /callbacks: default
  - override /trainer: default
  - override /logger: tensorboard

# load_ckpt_path: [PATH_TO_YOUR_CKPT] # uncomment to load a checkpoint
seed: 12345

trainer:
  max_steps: 300000
  val_check_interval: 10000 # this controls BOTH checkpoint steps and validation
  accumulate_grad_batches: 1 # gradient accumulation

data:
  batch_size: 1 # should be the same as the number of GPUs
  num_workers: 16

model:
  _target_: models.simplefold.simplefold.SimpleFold
  ema_decay: 0.999
  clip_grad_norm_val: 2.0
  esm_model: "esm2_3B"
  lddt_weight_schedule: False # set to True in finetuning phase

  scheduler:
    _target_: onescience.utils.simplefold.lr_scheduler.LinearWarmup
    _partial_: true
    min_lr: 1e-6
    max_lr: ${model.optimizer.lr}
    warmup_steps: 5000

  architecture:
    esm_model: ${model.esm_model}

  processor:
    scale: 16.0
    ref_scale: 5.0
    multiplicity: 16 # number of copies per GPU