Spaces:
Sleeping
Sleeping
| """ | |
| Train the VGAE model on the pre-built PyG graph. | |
| Run from the repository root: | |
| python -m src.gnn.trainer | |
| """ | |
| import os | |
| import torch | |
| import torch.nn.functional as F | |
| from torch_geometric.utils import negative_sampling | |
| from src.gnn.model import get_model | |
| def train_gnn(pyg_data_path: str = "./storage_graph/pyg_data.pt", epochs: int = 100): | |
| print(f"Loading PyG data from {pyg_data_path}…") | |
| if not os.path.exists(pyg_data_path): | |
| print(f"File not found: {pyg_data_path}") | |
| return | |
| # Allowlist PyG classes for safe loading under PyTorch >= 2.6. | |
| try: | |
| from torch_geometric.data import Data | |
| from torch_geometric.data.data import DataEdgeAttr | |
| torch.serialization.add_safe_globals([Data, DataEdgeAttr]) | |
| except (ImportError, AttributeError): | |
| pass | |
| checkpoint = torch.load(pyg_data_path, weights_only=False) | |
| data = checkpoint["pyg_data"] | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| data = data.to(device) | |
| in_channels = data.x.size(1) | |
| print(f"Initialising VGAE: in={in_channels}, hidden=256, out=128 — device={device}") | |
| model = get_model(in_channels=in_channels, hidden_channels=256, out_channels=128).to(device) | |
| optimizer = torch.optim.Adam(model.parameters(), lr=0.01) | |
| print("Training with link-prediction objective (L_recon + L_KL)…") | |
| model.train() | |
| for epoch in range(epochs): | |
| optimizer.zero_grad() | |
| z = model.encode(data.x, data.edge_index) | |
| # Reconstruction loss on positive edges. | |
| pos_loss = -torch.log(model.decode(z, data.edge_index, sigmoid=True) + 1e-15).mean() | |
| # Reconstruction loss on sampled negative edges (|E⁻| = |E⁺|). | |
| neg_edge_index = negative_sampling( | |
| edge_index=data.edge_index, | |
| num_nodes=data.num_nodes, | |
| num_neg_samples=data.edge_index.size(1), | |
| ) | |
| neg_loss = -torch.log(1 - model.decode(z, neg_edge_index, sigmoid=True) + 1e-15).mean() | |
| # KL regularisation term — requires encode() to have been called above. | |
| kl = model.kl_loss() | |
| loss = pos_loss + neg_loss + kl | |
| loss.backward() | |
| optimizer.step() | |
| if epoch % 10 == 0: | |
| print(f"Epoch {epoch:03d}: loss={loss.item():.4f} " | |
| f"(recon={pos_loss.item()+neg_loss.item():.4f}, kl={kl.item():.4f})") | |
| print("Training complete.") | |
| with torch.no_grad(): | |
| model.eval() | |
| z_final = model.encode(data.x, data.edge_index).cpu() | |
| print("Saving model weights and structural embeddings…") | |
| checkpoint["structural_embeddings"] = z_final | |
| torch.save(checkpoint, pyg_data_path) | |
| torch.save(model.state_dict(), "./storage_graph/gnn_model.pth") | |
| print("Done.") | |
| if __name__ == "__main__": | |
| train_gnn() | |