| 43 | |
| 44 | |
| 45 | class CrossScaleSampler(Sampler): |
| 46 | def __init__(self, *args, **kwargs): |
| 47 | super(CrossScaleSampler, self).__init__(*args, **kwargs) |
| 48 | |
| 49 | def sample(self, anchor, sample, batch=None, neg_sample=None, use_gpu=True, *args, **kwargs): |
| 50 | num_graphs = anchor.shape[0] # M |
| 51 | num_nodes = sample.shape[0] # N |
| 52 | device = sample.device |
| 53 | |
| 54 | if neg_sample is not None: |
| 55 | assert num_graphs == 1 # only one graph, explicit negative samples are needed |
| 56 | assert sample.shape == neg_sample.shape |
| 57 | pos_mask1 = torch.ones((num_graphs, num_nodes), dtype=torch.float32, device=device) |
| 58 | pos_mask0 = torch.zeros((num_graphs, num_nodes), dtype=torch.float32, device=device) |
| 59 | pos_mask = torch.cat([pos_mask1, pos_mask0], dim=1) # M * 2N |
| 60 | sample = torch.cat([sample, neg_sample], dim=0) # 2N * K |
| 61 | else: |
| 62 | assert batch is not None |
| 63 | if use_gpu: |
| 64 | ones = torch.eye(num_nodes, dtype=torch.float32, device=device) # N * N |
| 65 | pos_mask = scatter(ones, batch, dim=0, reduce='sum') # M * N |
| 66 | else: |
| 67 | pos_mask = torch.zeros((num_graphs, num_nodes), dtype=torch.float32).to(device) |
| 68 | for node_idx, graph_idx in enumerate(batch): |
| 69 | pos_mask[graph_idx][node_idx] = 1. # M * N |
| 70 | |
| 71 | neg_mask = 1. - pos_mask |
| 72 | return anchor, sample, pos_mask, neg_mask |
| 73 | |
| 74 | |
| 75 | def get_sampler(mode: str, intraview_negs: bool) -> Sampler: |