| 37 | |
| 38 | |
| 39 | class GConv(torch.nn.Module): |
| 40 | def __init__(self, input_dim, hidden_dim, num_layers, dropout=0.2, |
| 41 | encoder_norm='batch', projector_norm='batch'): |
| 42 | super(GConv, self).__init__() |
| 43 | self.activation = torch.nn.PReLU() |
| 44 | self.dropout = dropout |
| 45 | |
| 46 | self.layers = torch.nn.ModuleList() |
| 47 | self.layers.append(make_gin_conv(input_dim, hidden_dim)) |
| 48 | for _ in range(num_layers - 1): |
| 49 | self.layers.append(make_gin_conv(hidden_dim, hidden_dim)) |
| 50 | |
| 51 | self.batch_norm = Normalize(hidden_dim, norm=encoder_norm) |
| 52 | self.projection_head = torch.nn.Sequential( |
| 53 | torch.nn.Linear(hidden_dim, hidden_dim), |
| 54 | Normalize(hidden_dim, norm=projector_norm), |
| 55 | torch.nn.PReLU(), |
| 56 | torch.nn.Dropout(dropout)) |
| 57 | |
| 58 | def forward(self, x, edge_index, edge_weight=None): |
| 59 | z = x |
| 60 | for conv in self.layers: |
| 61 | z = conv(z, edge_index, edge_weight) |
| 62 | z = self.activation(z) |
| 63 | z = F.dropout(z, p=self.dropout, training=self.training) |
| 64 | z = self.batch_norm(z) |
| 65 | return z, self.projection_head(z) |
| 66 | |
| 67 | |
| 68 | class Encoder(torch.nn.Module): |