| 17 | |
| 18 | |
| 19 | class GConv(nn.Module): |
| 20 | def __init__(self, input_dim, hidden_dim, activation, num_layers): |
| 21 | super(GConv, self).__init__() |
| 22 | self.activation = activation() |
| 23 | self.layers = nn.ModuleList() |
| 24 | self.batch_norms = nn.ModuleList() |
| 25 | for i in range(num_layers): |
| 26 | if i == 0: |
| 27 | self.layers.append(make_gin_conv(input_dim, hidden_dim)) |
| 28 | else: |
| 29 | self.layers.append(make_gin_conv(hidden_dim, hidden_dim)) |
| 30 | self.batch_norms.append(nn.BatchNorm1d(hidden_dim)) |
| 31 | |
| 32 | def forward(self, x, edge_index, batch): |
| 33 | z = x |
| 34 | zs = [] |
| 35 | for conv, bn in zip(self.layers, self.batch_norms): |
| 36 | z = conv(z, edge_index) |
| 37 | z = self.activation(z) |
| 38 | z = bn(z) |
| 39 | zs.append(z) |
| 40 | gs = [global_add_pool(z, batch) for z in zs] |
| 41 | z, g = [torch.cat(x, dim=1) for x in [zs, gs]] |
| 42 | return z, g |
| 43 | |
| 44 | |
| 45 | class FC(nn.Module): |