| 34 | |
| 35 | |
| 36 | class Encoder(torch.nn.Module): |
| 37 | def __init__(self, encoder1, encoder2, augmentor, hidden_dim): |
| 38 | super(Encoder, self).__init__() |
| 39 | self.encoder1 = encoder1 |
| 40 | self.encoder2 = encoder2 |
| 41 | self.augmentor = augmentor |
| 42 | self.project = torch.nn.Linear(hidden_dim, hidden_dim) |
| 43 | uniform(hidden_dim, self.project.weight) |
| 44 | |
| 45 | @staticmethod |
| 46 | def corruption(x, edge_index, edge_weight): |
| 47 | return x[torch.randperm(x.size(0))], edge_index, edge_weight |
| 48 | |
| 49 | def forward(self, x, edge_index, edge_weight=None): |
| 50 | aug1, aug2 = self.augmentor |
| 51 | x1, edge_index1, edge_weight1 = aug1(x, edge_index, edge_weight) |
| 52 | x2, edge_index2, edge_weight2 = aug2(x, edge_index, edge_weight) |
| 53 | z1 = self.encoder1(x1, edge_index1, edge_weight1) |
| 54 | z2 = self.encoder2(x2, edge_index2, edge_weight2) |
| 55 | g1 = self.project(torch.sigmoid(z1.mean(dim=0, keepdim=True))) |
| 56 | g2 = self.project(torch.sigmoid(z2.mean(dim=0, keepdim=True))) |
| 57 | z1n = self.encoder1(*self.corruption(x1, edge_index1, edge_weight1)) |
| 58 | z2n = self.encoder2(*self.corruption(x2, edge_index2, edge_weight2)) |
| 59 | return z1, z2, g1, g2, z1n, z2n |
| 60 | |
| 61 | |
| 62 | def train(encoder_model, contrast_model, data, optimizer): |