Updated all code

model.py

@@ -94,7 +94,7 @@ class ProteinTransformer(nn.Module):
         super().__init__()
         self.d_model = d_model
         self.embedding = nn.Embedding(vocab_size, d_model)
-        self.pos_encoder = PositionalEncoding(d_model, dropout)
+        self.pos_encoder = PositionalEncoding(d_model, dropout=dropout)
         encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=h, batch_first=True)
         self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=N)
 

optuna_train.py

@@ -1,24 +1,34 @@
+import optuna
 import torch
 import torch.nn as nn
 import pandas as pd
-import optuna
-from torch.nn.functional import dropout
-from torch.utils.data import random_split
+import random
+import numpy as np
 from torch_geometric.loader import DataLoader
+from torch.utils.data import random_split
 from dataset import BindingDataset
 from model import BindingAffinityModel
-from tqdm import tqdm
-import sys
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-EPOCHS_PER_TRIAL = 10
+N_TRIALS = 20
+EPOCHS_PER_TRIAL = 15
+
+def set_seed(seed=42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    return torch.Generator().manual_seed(seed)
 
 dataframe = pd.read_csv('pdbbind_refined_dataset.csv')
 dataframe.dropna(inplace=True)
 dataset = BindingDataset(dataframe)
+
+gen = set_seed(42)
+
 train_size = int(0.8 * len(dataset))
 test_size = len(dataset) - train_size
-train_dataset, test_dataset = random_split(dataset, [train_size, test_size])
+train_dataset, test_dataset = random_split(dataset, [train_size, test_size], generator=gen)
 num_features = train_dataset[0].x.shape[1]
 
 def train(model, loader, optimizer, criterion):
@@ -45,21 +55,31 @@ def test(model, loader, criterion):
 
 
 def objective(trial):
+    # Architecture
+    hidden_dim = trial.suggest_categorical("hidden_dim", [64, 128, 256])
+    gat_heads = trial.suggest_categorical("gat_heads", [2, 4, 8])
+    dropout = trial.suggest_float("dropout", 0.1, 0.5)
+
+    # Learning
+
     lr = trial.suggest_float("lr", 1e-5, 1e-2, log=True) # Learning rate from 0.00001 to 0.01
     weight_decay = trial.suggest_float("weight_decay", 1e-6, 1e-3, log=True) # Weight decay from 0.000001 to 0.001
+    batch_size = trial.suggest_categorical("batch_size", [16, 32, 64])
 
-    model = BindingAffinityModel(num_node_features=num_features, hidden_channels_gnn=128).to(DEVICE)
+    model = BindingAffinityModel(num_node_features=num_features, hidden_channels=hidden_dim, gat_heads=gat_heads, dropout=dropout).to(DEVICE)
 
     optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
     criterion = nn.MSELoss()
 
-    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
-    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
 
     for epoch in range(EPOCHS_PER_TRIAL):
         train(model, train_loader, optimizer, criterion)
         val_loss = test(model, test_loader, criterion)
 
+        print(f"Trial {trial.number} | Epoch {epoch + 1}/{EPOCHS_PER_TRIAL} | Val Loss: {val_loss:.4f}")
+
         trial.report(val_loss, epoch)
         if trial.should_prune():
             raise optuna.exceptions.TrialPruned()
@@ -67,10 +87,17 @@ def objective(trial):
 
 
 if __name__ == "__main__":
-    study = optuna.create_study(direction="minimize")
+    storage_name = "sqlite:///db.sqlite3"
+    study = optuna.create_study(
+        direction="minimize",
+        pruner=optuna.pruners.MedianPruner(),
+        storage=storage_name,
+        study_name="binding_prediction_optimization",
+        load_if_exists=True
+    )
     print("Start hyperparameter optimization...")
 
-    study.optimize(objective, n_trials=10)
+    study.optimize(objective, n_trials=N_TRIALS)
     print("\n--- Optimization Finished ---")
     print("Best parameters found: ", study.best_params)
     print("Best Test MSE: ", study.best_value)

train.py

@@ -11,12 +11,20 @@ from tqdm import tqdm
 from torch.utils.tensorboard import SummaryWriter
 import numpy as np
 from datetime import datetime
+import os
+
+
+BATCH_SIZE = 16
+LR = 0.00064
+WEIGHT_DECAY = 7.06e-6
+EPOCS = 100
+DROPOUT = 0.325
+GAT_HEADS = 2
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-BATCH_SIZE = 32
-LR = 0.0005
-EPOCS = 30
 LOG_DIR = f"runs/experiment_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+TOP_K = 3
+SAVES_DIR = LOG_DIR + "/models"
 
 def set_seed(seed=42):
     random.seed(seed)
@@ -68,7 +76,11 @@ def evaluate(epoch, model, loader, criterion, writer):
 def main():
     gen = set_seed(42)
     writer = SummaryWriter(LOG_DIR)
+
+    if not os.path.exists(SAVES_DIR):
+        os.makedirs(SAVES_DIR)
     print(f"Logging to {LOG_DIR}...")
+    print(f"Model saves to {SAVES_DIR}...")
     # Load dataset
     dataframe = pd.read_csv('pdbbind_refined_dataset.csv')
     dataframe.dropna(inplace=True)
@@ -90,26 +102,47 @@ def main():
     num_features = train_dataset[0].x.shape[1]
     print("Number of node features:", num_features)
 
-    model = BindingAffinityModel(num_node_features=num_features, hidden_channels_gnn=128).to(DEVICE)
-    optimizer = torch.optim.Adam(model.parameters(), lr=0.0005, weight_decay=1e-4)
+    model = BindingAffinityModel(
+        num_node_features=num_features,
+        hidden_channels=256,
+        gat_heads=GAT_HEADS,
+        dropout=DROPOUT
+    ).to(DEVICE)
+    optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
     criterion = nn.MSELoss()
 
-    best_test_loss = float('inf')
+    top_models = []
 
     print(f"Starting training on {DEVICE}")
-    for epoch in range(1, EPOCS):
+    for epoch in range(1, EPOCS + 1):
         train_loss = train_epoch(epoch, model, train_loader, optimizer, criterion, writer)
         test_loss = evaluate(epoch, model, test_loader, criterion, writer)
 
-
         print(f'Epoch {epoch:02d}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}')
-        if test_loss < best_test_loss:
-            best_test_loss = test_loss
-            torch.save(model.state_dict(), f'best_model_gat.pth')
-            print(f'Best model saved with Test Loss MSE: {best_test_loss:.4f}')
+
+        filename = f"{SAVES_DIR}/model_ep{epoch:03d}_mse{test_loss:.4f}.pth"
+
+        torch.save(model.state_dict(), filename)
+        top_models.append({'loss': test_loss, 'path': filename, 'epoch': epoch})
+
+        top_models.sort(key=lambda x: x['loss'])
+
+        if len(top_models) > TOP_K:
+            worst_model = top_models.pop()
+            os.remove(worst_model['path'])
+
+        if any(m['epoch'] == epoch for m in top_models):
+            rank = [m['epoch'] for m in top_models].index(epoch) + 1
+            print(f'-- Model saved (Rank: {rank})')
+        else:
+            print("")
+
 
     writer.close()
     print("Training finished.")
+    print("Top models saved:")
+    for i, m in enumerate(top_models):
+        print(f"{i + 1}. {m['path']} (MSE: {m['loss']:.4f})")
 
 
 if __name__ == "__main__":