"sgl-kernel/python/vscode:/vscode.git/clone" did not exist on "8bf6d7f406145cc46b5dcfa900e7710417fe2a8a"
negative_sampler.py 728 Bytes
Newer Older
1
2
3
4
import torch as th
import dgl

class NegativeSampler(object):
5
6
7
8
    def __init__(self, g, k, neg_share=False, device=None):
        if device is None:
            device = g.device
        self.weights = g.in_degrees().float().to(device) ** 0.75
9
10
11
12
13
14
15
16
17
18
19
20
21
        self.k = k
        self.neg_share = neg_share

    def __call__(self, g, eids):
        src, _ = g.find_edges(eids)
        n = len(src)
        if self.neg_share and n % self.k == 0:
            dst = self.weights.multinomial(n, replacement=True)
            dst = dst.view(-1, 1, self.k).expand(-1, self.k, -1).flatten()
        else:
            dst = self.weights.multinomial(n*self.k, replacement=True)
        src = src.repeat_interleave(self.k)
        return src, dst