Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch_geometric.nn import SAGEConv | |
| class VGAEEncoder(nn.Module): | |
| """ | |
| GraphSAGE-based variational encoder for VGAE. | |
| SAGEConv uses mean aggregation, which is degree-agnostic and more stable | |
| than GCNConv's degree-normalised aggregation on sparse knowledge graphs. | |
| Outputs (mu, log_std) for the reparameterisation trick. | |
| """ | |
| def __init__(self, in_channels: int, hidden_channels: int, out_channels: int, dropout: float = 0.1): | |
| super().__init__() | |
| self.conv1 = SAGEConv(in_channels, hidden_channels) | |
| self.conv_mu = SAGEConv(hidden_channels, out_channels) | |
| self.conv_logstd = SAGEConv(hidden_channels, out_channels) | |
| self.dropout = dropout | |
| self._init_weights() | |
| def _init_weights(self): | |
| # Xavier uniform initialisation to control activation magnitude from epoch 0. | |
| for conv in [self.conv1, self.conv_mu, self.conv_logstd]: | |
| nn.init.xavier_uniform_(conv.lin_l.weight) | |
| nn.init.xavier_uniform_(conv.lin_r.weight) | |
| def forward(self, x, edge_index): | |
| x = self.conv1(x, edge_index) | |
| x = F.relu(x) | |
| x = F.dropout(x, p=self.dropout, training=self.training) | |
| mu = self.conv_mu(x, edge_index) | |
| # Clamp both directions: prevents underflow (exp(-10)≈0) and overflow (exp(10)≈22k). | |
| log_std = self.conv_logstd(x, edge_index).clamp(min=-10, max=10) | |
| return mu, log_std | |
| class VGAEModel(nn.Module): | |
| def __init__(self, in_channels: int, hidden_channels: int, out_channels: int, dropout: float = 0.1): | |
| super().__init__() | |
| self.encoder = VGAEEncoder(in_channels, hidden_channels, out_channels, dropout) | |
| def encode(self, x, edge_index): | |
| mu, log_std = self.encoder(x, edge_index) | |
| # At training time use the reparameterisation trick; at inference use the | |
| # deterministic mean mu for reproducible retrieval scores. | |
| if self.training: | |
| std = torch.exp(log_std) | |
| z = mu + torch.randn_like(std) * std | |
| else: | |
| z = mu | |
| self._mu = mu | |
| self._log_std = log_std | |
| return z | |
| def decode(self, z, edge_index, sigmoid=True): | |
| value = (z[edge_index[0]] * z[edge_index[1]]).sum(dim=1) | |
| return torch.sigmoid(value) if sigmoid else value | |
| def kl_loss(self): | |
| """KL( N(mu, std) || N(0, 1) ) — call after encode().""" | |
| return -0.5 * torch.mean( | |
| 1 + 2 * self._log_std - self._mu.pow(2) - (2 * self._log_std).exp() | |
| ) | |
| def forward(self, x, edge_index): | |
| return self.encode(x, edge_index) | |
| def get_model(in_channels: int = 384, hidden_channels: int = 256, out_channels: int = 128) -> VGAEModel: | |
| return VGAEModel(in_channels, hidden_channels, out_channels) | |