Spaces:
Running
Running
| import torch | |
| from torch_geometric.loader import DataLoader | |
| import torch_geometric | |
| import numpy as np | |
| import json | |
| import os | |
| from dotenv import load_dotenv | |
| from src.model import GIN | |
| from src.preprocess import get_graph_datasets | |
| from src.train_evaluate import train_model, evaluate, compute_roc_auc_avg_and_per_class | |
| from src.seed import set_seed | |
| def train(config): | |
| SEED=config["seed"] | |
| set_seed(SEED) | |
| best_model_path = "./checkpoints/model.pt" | |
| # get dataloaders | |
| print("Loading Datasets...") | |
| torch_geometric.seed_everything(SEED) | |
| token = os.getenv("TOKEN") | |
| train_dataset, val_dataset = get_graph_datasets(token) | |
| train_loader = DataLoader(train_dataset, batch_size=config["batch_size"], shuffle=True) | |
| val_loader = DataLoader(val_dataset, batch_size=config["batch_size"]) | |
| # initialize | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = GIN(num_features=9, num_classes=12, dropout=config["dropout"], hidden_dim=config["hidden_dim"], num_layers=config["num_layers"], add_or_mean=config["add_or_mean"]).to(device) | |
| optimizer = torch.optim.Adam(model.parameters(), lr=config["lr"]) | |
| # training loop | |
| best_mean_auc = -float("inf") | |
| best_mean_epoch = 0 | |
| aucs = [] | |
| window_size = config["window_size"] | |
| epoch_checkpoints = {} | |
| print("Starting Training...") | |
| for epoch in range(0, config["max_epochs"]): | |
| train_loss = train_model(model, train_loader, optimizer, device) | |
| val_loss = evaluate(model, val_loader, device) | |
| val_auc_per_class, val_auc_avg = compute_roc_auc_avg_and_per_class(model, val_loader, device) | |
| aucs.append(val_auc_avg) | |
| # log | |
| if epoch % 10 == 0: | |
| print(f"Epoch {epoch:03d} | " | |
| f"Train Loss: {train_loss:.4f} | " | |
| f"Val Loss: {val_loss:.4f} | " | |
| f"Val ROC-AUC: {val_auc_avg:.4f}") | |
| # store model parameters for this epoch in cache (on CPU to save GPU memory) | |
| epoch_checkpoints[epoch] = {k: v.cpu() for k, v in model.state_dict().items()} | |
| # keep cache size limited | |
| if len(epoch_checkpoints) > window_size + 2: | |
| oldest = min(epoch_checkpoints.keys()) | |
| del epoch_checkpoints[oldest] | |
| # once we have enough epochs, compute rolling mean | |
| if len(aucs) >= window_size: | |
| current_window = aucs[-window_size:] | |
| current_mean_auc = np.mean(current_window) | |
| middle_epoch = epoch - window_size // 2 | |
| # check if current mean beats the best so far | |
| if current_mean_auc > best_mean_auc: | |
| best_mean_auc = current_mean_auc | |
| best_mean_epoch = middle_epoch | |
| # save only the best middle model | |
| if middle_epoch in epoch_checkpoints: | |
| torch.save(epoch_checkpoints[middle_epoch], best_model_path) | |
| print(f"🟢 New best mean AUC = {best_mean_auc:.4f} " | |
| f"(center epoch {best_mean_epoch}) — model saved!") | |
| # early stopping based on best mean epoch | |
| if epoch - best_mean_epoch >= config["patience"]: | |
| print(f"⛔ Early stopping at epoch {epoch}. " | |
| f"Best mean AUC = {best_mean_auc:.4f} (center epoch {best_mean_epoch})") | |
| break | |
| print("best_smoothed_val_auc" + str(best_mean_auc) + ", best_middle_epoch" + str(best_mean_epoch)) | |
| if __name__ == "__main__": | |
| with open("./config/config.json", "r") as f: | |
| config = json.load(f) | |
| load_dotenv() | |
| train(config) |