File size: 1,367 Bytes
b142f2b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
{
  "seed": 42,
  "wandb": true,
  "truncation_seq_length": 512,
  "batch_size": 16,
  "datamodule": {
    "batch_size": 64,
    "num_workers": 32,
    "_target_": "foldmae.data.TEDLightningDataset",
    "root": "./datasets/afdb_FS_plddt80",
    "dataset_name": "ted"
  },
  "trainer": {
    "_target_": "pytorch_lightning.Trainer",
    "accelerator": "auto",
    "max_steps": 18300,
    "strategy": "ddp_find_unused_parameters_true",
    "devices": "auto",
    "default_root_dir": "${logs.path}",
    "num_sanity_val_steps": 0,
    "accumulate_grad_batches": 16
  },
  "model": {
    "name": "esm2_t33_650M_UR50D",
    "num_classes": 965,
    "avg_pool": false
  },
  "train": {
    "optimizer": {
      "_target_": "torch.optim.AdamW",
      "lr": 0.0016,
      "weight_decay": 0.1,
      "betas": [
        0.9,
        0.95
      ]
    },
    "lr_scheduler": {
      "_target_": "foldmae.lr_schedulers.get_cosine_schedule_with_warmup",
      "warmup_steps": 1830,
      "max_steps": "${trainer.max_steps}"
    },
    "loss": {
      "_target_": "torch.nn.CrossEntropyLoss"
    },
    "llrd": 0.8,
    "ckpt_path": "logs/esm2/finetune/ted/42/runs/2026-01-20_07-59-12/model.ckpt"
  },
  "logs": {
    "prefix": "logs/esm2/finetune/${datamodule.dataset_name}/${seed}",
    "path": "${logs.prefix}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S}"
  },
  "_model_class": "esm2"
}