(self, h1_pred=None, h2_pred=None, h1_target=None, h2_target=None,
g1_pred=None, g2_pred=None, g1_target=None, g2_target=None,
batch=None, extra_pos_mask=None)
| 80 | self.sampler = get_sampler(mode, intraview_negs=False) |
| 81 | |
| 82 | def forward(self, h1_pred=None, h2_pred=None, h1_target=None, h2_target=None, |
| 83 | g1_pred=None, g2_pred=None, g1_target=None, g2_target=None, |
| 84 | batch=None, extra_pos_mask=None): |
| 85 | if self.mode == 'L2L': |
| 86 | assert all(v is not None for v in [h1_pred, h2_pred, h1_target, h2_target]) |
| 87 | anchor1, sample1, pos_mask1, _ = self.sampler(anchor=h1_target, sample=h2_pred) |
| 88 | anchor2, sample2, pos_mask2, _ = self.sampler(anchor=h2_target, sample=h1_pred) |
| 89 | elif self.mode == 'G2G': |
| 90 | assert all(v is not None for v in [g1_pred, g2_pred, g1_target, g2_target]) |
| 91 | anchor1, sample1, pos_mask1, _ = self.sampler(anchor=g1_target, sample=g2_pred) |
| 92 | anchor2, sample2, pos_mask2, _ = self.sampler(anchor=g2_target, sample=g1_pred) |
| 93 | else: |
| 94 | assert all(v is not None for v in [h1_pred, h2_pred, g1_target, g2_target]) |
| 95 | if batch is None or batch.max().item() + 1 <= 1: # single graph |
| 96 | pos_mask1 = pos_mask2 = torch.ones([1, h1_pred.shape[0]], device=h1_pred.device) |
| 97 | anchor1, sample1 = g1_target, h2_pred |
| 98 | anchor2, sample2 = g2_target, h1_pred |
| 99 | else: |
| 100 | anchor1, sample1, pos_mask1, _ = self.sampler(anchor=g1_target, sample=h2_pred, batch=batch) |
| 101 | anchor2, sample2, pos_mask2, _ = self.sampler(anchor=g2_target, sample=h1_pred, batch=batch) |
| 102 | |
| 103 | pos_mask1, _ = add_extra_mask(pos_mask1, extra_pos_mask=extra_pos_mask) |
| 104 | pos_mask2, _ = add_extra_mask(pos_mask2, extra_pos_mask=extra_pos_mask) |
| 105 | l1 = self.loss(anchor=anchor1, sample=sample1, pos_mask=pos_mask1) |
| 106 | l2 = self.loss(anchor=anchor2, sample=sample2, pos_mask=pos_mask2) |
| 107 | |
| 108 | return (l1 + l2) * 0.5 |
| 109 | |
| 110 | |
| 111 | class WithinEmbedContrast(torch.nn.Module): |
nothing calls this directly
no test coverage detected