| 54 | |
| 55 | |
| 56 | def train(encoder_model, contrast_model, data, optimizer): |
| 57 | encoder_model.train() |
| 58 | optimizer.zero_grad() |
| 59 | z, z1, z2 = encoder_model(data.x, data.edge_index, data.edge_attr) |
| 60 | h1, h2 = [encoder_model.project(x) for x in [z1, z2]] |
| 61 | |
| 62 | # compute extra pos and neg masks for semi-supervised learning |
| 63 | extra_pos_mask = torch.eq(data.y, data.y.unsqueeze(dim=1)).to('cuda') |
| 64 | # construct extra supervision signals for only training samples |
| 65 | extra_pos_mask[~data.train_mask][:, ~data.train_mask] = False |
| 66 | extra_pos_mask.fill_diagonal_(False) |
| 67 | # pos_mask: [N, 2N] for both inter-view and intra-view samples |
| 68 | extra_pos_mask = torch.cat([extra_pos_mask, extra_pos_mask], dim=1).to('cuda') |
| 69 | # fill interview positives only; pos_mask for intraview samples should have zeros in diagonal |
| 70 | extra_pos_mask.fill_diagonal_(True) |
| 71 | |
| 72 | extra_neg_mask = torch.ne(data.y, data.y.unsqueeze(dim=1)).to('cuda') |
| 73 | extra_neg_mask[~data.train_mask][:, ~data.train_mask] = True |
| 74 | extra_neg_mask.fill_diagonal_(False) |
| 75 | extra_neg_mask = torch.cat([extra_neg_mask, extra_neg_mask], dim=1).to('cuda') |
| 76 | |
| 77 | loss = contrast_model(h1=h1, h2=h2, extra_pos_mask=extra_pos_mask, extra_neg_mask=extra_neg_mask) |
| 78 | loss.backward() |
| 79 | optimizer.step() |
| 80 | return loss.item() |
| 81 | |
| 82 | |
| 83 | def test(encoder_model, data): |