| |
|
|
| class GNNDiracRRF(nn.Module): |
| def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int, |
| alpha_attn: float = 1.0, dropout: float = 0.1): |
| super().__init__() |
| self.z_dim = z_dim |
| self.layers = nn.ModuleList() |
| self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn)) |
| for _ in range(num_layers - 2): |
| self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn)) |
| self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn)) |
| self.dropout = nn.Dropout(dropout) |
|
|
| def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor: |
| h = x |
| for i, layer in enumerate(self.layers): |
| h = layer(h, edge_index, z) |
| if i < len(self.layers) - 1: |
| h = F.gelu(h) |
| h = self.dropout(h) |
| return h |
