| 14 | |
| 15 | |
| 16 | class GConv(nn.Module): |
| 17 | def __init__(self, input_dim, hidden_dim, num_layers): |
| 18 | super(GConv, self).__init__() |
| 19 | self.layers = torch.nn.ModuleList() |
| 20 | self.activations = torch.nn.ModuleList() |
| 21 | for i in range(num_layers): |
| 22 | if i == 0: |
| 23 | self.layers.append(SAGEConv(input_dim, hidden_dim)) |
| 24 | else: |
| 25 | self.layers.append(SAGEConv(hidden_dim, hidden_dim)) |
| 26 | self.activations.append(nn.PReLU(hidden_dim)) |
| 27 | |
| 28 | def forward(self, x, adjs): |
| 29 | for i, (edge_index, _, size) in enumerate(adjs): |
| 30 | x_target = x[:size[1]] |
| 31 | x = self.layers[i]((x, x_target), edge_index) |
| 32 | x = self.activations[i](x) |
| 33 | return x |
| 34 | |
| 35 | |
| 36 | class Encoder(torch.nn.Module): |