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