(self, h1=None, h2=None, g1=None, g2=None, batch=None, h3=None, h4=None,
extra_pos_mask=None, extra_neg_mask=None)
| 45 | self.kwargs = kwargs |
| 46 | |
| 47 | def forward(self, h1=None, h2=None, g1=None, g2=None, batch=None, h3=None, h4=None, |
| 48 | extra_pos_mask=None, extra_neg_mask=None): |
| 49 | if self.mode == 'L2L': |
| 50 | assert h1 is not None and h2 is not None |
| 51 | anchor1, sample1, pos_mask1, neg_mask1 = self.sampler(anchor=h1, sample=h2) |
| 52 | anchor2, sample2, pos_mask2, neg_mask2 = self.sampler(anchor=h2, sample=h1) |
| 53 | elif self.mode == 'G2G': |
| 54 | assert g1 is not None and g2 is not None |
| 55 | anchor1, sample1, pos_mask1, neg_mask1 = self.sampler(anchor=g1, sample=g2) |
| 56 | anchor2, sample2, pos_mask2, neg_mask2 = self.sampler(anchor=g2, sample=g1) |
| 57 | else: # global-to-local |
| 58 | if batch is None or batch.max().item() + 1 <= 1: # single graph |
| 59 | assert all(v is not None for v in [h1, h2, g1, g2, h3, h4]) |
| 60 | anchor1, sample1, pos_mask1, neg_mask1 = self.sampler(anchor=g1, sample=h2, neg_sample=h4) |
| 61 | anchor2, sample2, pos_mask2, neg_mask2 = self.sampler(anchor=g2, sample=h1, neg_sample=h3) |
| 62 | else: # multiple graphs |
| 63 | assert all(v is not None for v in [h1, h2, g1, g2, batch]) |
| 64 | anchor1, sample1, pos_mask1, neg_mask1 = self.sampler(anchor=g1, sample=h2, batch=batch) |
| 65 | anchor2, sample2, pos_mask2, neg_mask2 = self.sampler(anchor=g2, sample=h1, batch=batch) |
| 66 | |
| 67 | pos_mask1, neg_mask1 = add_extra_mask(pos_mask1, neg_mask1, extra_pos_mask, extra_neg_mask) |
| 68 | pos_mask2, neg_mask2 = add_extra_mask(pos_mask2, neg_mask2, extra_pos_mask, extra_neg_mask) |
| 69 | l1 = self.loss(anchor=anchor1, sample=sample1, pos_mask=pos_mask1, neg_mask=neg_mask1, **self.kwargs) |
| 70 | l2 = self.loss(anchor=anchor2, sample=sample2, pos_mask=pos_mask2, neg_mask=neg_mask2, **self.kwargs) |
| 71 | |
| 72 | return (l1 + l2) * 0.5 |
| 73 | |
| 74 | |
| 75 | class BootstrapContrast(torch.nn.Module): |
nothing calls this directly
no test coverage detected