foldtree2/config_notebook_example.yaml at main · DessimozLab/foldtree2 · GitHub

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
# Example Configuration File for test_monodecoders_config.ipynb
# This file demonstrates how to configure all hyperparameters from a YAML file
# Place this file in the project root and reference it in the notebook's config_file variable

# ============================================================================
# DATA PARAMETERS
# ============================================================================
datadir: '../../datasets/foldtree2/'
dataset_path: 'structs_train_final.h5'
aapropcsv: './foldtree2/config/aaindex1.csv'

# ============================================================================
# MODEL ARCHITECTURE PARAMETERS
# ============================================================================
# Alphabet/Embedding parameters
num_embeddings: 30
embedding_dim: 128

# Network size
hidden_size: 150

# Encoder parameters
encoder_type: 'mk1_Encoder'
encoder_hidden_channels: [150, 150, 150]
encoder_nheads: 16
encoder_dropout: 0.005
encoder_flavor: 'transformer'
encoder_fftin: true
encoder_learn_positions: true
encoder_concat_positions: false

# Decoder parameters
use_monodecoder: true
use_muon_decoders: true

# ============================================================================
# TRAINING PARAMETERS
# ============================================================================
num_epochs: 300
batch_size: 10
gradient_accumulation_steps: 2
clip_grad: true
mask_plddt: true
plddt_threshold: 0.3
num_workers: 4

# Learning rate and scheduler
learning_rate: 0.00001  # 1e-5
scheduler_type: 'plateau'  # Options: 'plateau', 'linear', 'cosine', 'cosine_with_restarts', 'polynomial'
warmup_steps: 20
warmup_ratio: 0.05

# Optimizer parameters
use_muon: true
muon_lr: 0.02
adamw_lr: 0.0001  # 1e-4
weight_decay: 0.01

# ============================================================================
# LOSS WEIGHTS
# ============================================================================
edgeweight: 0.1
logitweight: 0.1
xweight: 0.1
fft2weight: 0.01
vqweight: 0.005
angles_weight: 0.1
ss_weight: 0.1

# Loss weight scheduler
use_weight_scheduler: true
loss_scheduler_type: 'linear'
loss_warmup_steps: 20

# ============================================================================
# COMMITMENT COST SCHEDULING
# ============================================================================
use_commitment_scheduling: true
commitment_cost_final: 0.9
commitment_warmup_steps: 1000
commitment_schedule: 'linear'  # Options: 'cosine', 'linear', 'none'
commitment_start: 0.5

# ============================================================================
# MIXED PRECISION TRAINING
# ============================================================================
use_mixed_precision: true

# ============================================================================
# REPRODUCIBILITY
# ============================================================================
random_seed: 0