| 34 | |
| 35 | |
| 36 | class Encoder(torch.nn.Module): |
| 37 | def __init__(self, encoder, hidden_dim): |
| 38 | super(Encoder, self).__init__() |
| 39 | self.encoder = encoder |
| 40 | self.project = torch.nn.Linear(hidden_dim, hidden_dim) |
| 41 | uniform(hidden_dim, self.project.weight) |
| 42 | |
| 43 | @staticmethod |
| 44 | def corruption(x, edge_index): |
| 45 | return x[torch.randperm(x.size(0))], edge_index |
| 46 | |
| 47 | def forward(self, x, edge_index): |
| 48 | z = self.encoder(x, edge_index) |
| 49 | g = self.project(torch.sigmoid(z.mean(dim=0, keepdim=True))) |
| 50 | zn = self.encoder(*self.corruption(x, edge_index)) |
| 51 | return z, g, zn |
| 52 | |
| 53 | |
| 54 | def train(encoder_model, contrast_model, data, dataloader, optimizer): |