| 15 | |
| 16 | |
| 17 | class SingleBranchContrast(torch.nn.Module): |
| 18 | def __init__(self, loss: Loss, mode: str, intraview_negs: bool = False, **kwargs): |
| 19 | super(SingleBranchContrast, self).__init__() |
| 20 | assert mode == 'G2L' # only global-local pairs allowed in single-branch contrastive learning |
| 21 | self.loss = loss |
| 22 | self.mode = mode |
| 23 | self.sampler = get_sampler(mode, intraview_negs=intraview_negs) |
| 24 | self.kwargs = kwargs |
| 25 | |
| 26 | def forward(self, h, g, batch=None, hn=None, extra_pos_mask=None, extra_neg_mask=None): |
| 27 | if batch is None: # for single-graph datasets |
| 28 | assert hn is not None |
| 29 | anchor, sample, pos_mask, neg_mask = self.sampler(anchor=g, sample=h, neg_sample=hn) |
| 30 | else: # for multi-graph datasets |
| 31 | assert batch is not None |
| 32 | anchor, sample, pos_mask, neg_mask = self.sampler(anchor=g, sample=h, batch=batch) |
| 33 | |
| 34 | pos_mask, neg_mask = add_extra_mask(pos_mask, neg_mask, extra_pos_mask, extra_neg_mask) |
| 35 | loss = self.loss(anchor=anchor, sample=sample, pos_mask=pos_mask, neg_mask=neg_mask, **self.kwargs) |
| 36 | return loss |
| 37 | |
| 38 | |
| 39 | class DualBranchContrast(torch.nn.Module): |