link.py 10.7 KB
Newer Older
1
2
import numpy as np
import torch
Mufei Li's avatar
Mufei Li committed
3
4
import torch.nn as nn
import torch.nn.functional as F
5
import dgl
Mufei Li's avatar
Mufei Li committed
6
7
from dgl.data.knowledge_graph import FB15k237Dataset
from dgl.dataloading import GraphDataLoader
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
from dgl.nn.pytorch import RelGraphConv
import tqdm

# for building training/testing graphs
def get_subset_g(g, mask, num_rels, bidirected=False):
    src, dst = g.edges()
    sub_src = src[mask]
    sub_dst = dst[mask]
    sub_rel = g.edata['etype'][mask]

    if bidirected:
        sub_src, sub_dst = torch.cat([sub_src, sub_dst]), torch.cat([sub_dst, sub_src])
        sub_rel = torch.cat([sub_rel, sub_rel + num_rels])

    sub_g = dgl.graph((sub_src, sub_dst), num_nodes=g.num_nodes())
    sub_g.edata[dgl.ETYPE] = sub_rel
    return sub_g

class GlobalUniform:
    def __init__(self, g, sample_size):
        self.sample_size = sample_size
        self.eids = np.arange(g.num_edges())

    def sample(self):
        return torch.from_numpy(np.random.choice(self.eids, self.sample_size))

class NegativeSampler:
    def __init__(self, k=10): # negative sampling rate = 10
        self.k = k

    def sample(self, pos_samples, num_nodes):
        batch_size = len(pos_samples)
        neg_batch_size = batch_size * self.k
        neg_samples = np.tile(pos_samples, (self.k, 1))

        values = np.random.randint(num_nodes, size=neg_batch_size)
        choices = np.random.uniform(size=neg_batch_size)
        subj = choices > 0.5
        obj = choices <= 0.5
        neg_samples[subj, 0] = values[subj]
        neg_samples[obj, 2] = values[obj]
        samples = np.concatenate((pos_samples, neg_samples))

        # binary labels indicating positive and negative samples
        labels = np.zeros(batch_size * (self.k + 1), dtype=np.float32)
        labels[:batch_size] = 1

        return torch.from_numpy(samples), torch.from_numpy(labels)

class SubgraphIterator:
    def __init__(self, g, num_rels, sample_size=30000, num_epochs=6000):
        self.g = g
        self.num_rels = num_rels
        self.sample_size = sample_size
        self.num_epochs = num_epochs
        self.pos_sampler = GlobalUniform(g, sample_size)
        self.neg_sampler = NegativeSampler()
Mufei Li's avatar
Mufei Li committed
65

66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
    def __len__(self):
        return self.num_epochs

    def __getitem__(self, i):
        eids = self.pos_sampler.sample()
        src, dst = self.g.find_edges(eids)
        src, dst = src.numpy(), dst.numpy()
        rel = self.g.edata[dgl.ETYPE][eids].numpy()

        # relabel nodes to have consecutive node IDs
        uniq_v, edges = np.unique((src, dst), return_inverse=True)
        num_nodes = len(uniq_v)
        # edges is the concatenation of src, dst with relabeled ID
        src, dst = np.reshape(edges, (2, -1))
        relabeled_data = np.stack((src, rel, dst)).transpose()

        samples, labels = self.neg_sampler.sample(relabeled_data, num_nodes)

        # use only half of the positive edges
        chosen_ids = np.random.choice(np.arange(self.sample_size),
                                      size=int(self.sample_size / 2),
                                      replace=False)
        src = src[chosen_ids]
        dst = dst[chosen_ids]
        rel = rel[chosen_ids]
        src, dst = np.concatenate((src, dst)), np.concatenate((dst, src))
        rel = np.concatenate((rel, rel + self.num_rels))
        sub_g = dgl.graph((src, dst), num_nodes=num_nodes)
        sub_g.edata[dgl.ETYPE] = torch.from_numpy(rel)
        sub_g.edata['norm'] = dgl.norm_by_dst(sub_g).unsqueeze(-1)
        uniq_v = torch.from_numpy(uniq_v).view(-1).long()

        return sub_g, uniq_v, samples, labels

class RGCN(nn.Module):
    def __init__(self, num_nodes, h_dim, num_rels):
        super().__init__()
        # two-layer RGCN
        self.emb = nn.Embedding(num_nodes, h_dim)
        self.conv1 = RelGraphConv(h_dim, h_dim, num_rels, regularizer='bdd',
                                  num_bases=100, self_loop=True)
        self.conv2 = RelGraphConv(h_dim, h_dim, num_rels, regularizer='bdd',
                                  num_bases=100, self_loop=True)
        self.dropout = nn.Dropout(0.2)

    def forward(self, g, nids):
        x = self.emb(nids)
        h = F.relu(self.conv1(g, x, g.edata[dgl.ETYPE], g.edata['norm']))
        h = self.dropout(h)
        h = self.conv2(g, h, g.edata[dgl.ETYPE], g.edata['norm'])
        return self.dropout(h)
Mufei Li's avatar
Mufei Li committed
117
118

class LinkPredict(nn.Module):
119
120
121
    def __init__(self, num_nodes, num_rels, h_dim = 500, reg_param=0.01):
        super().__init__()
        self.rgcn = RGCN(num_nodes, h_dim, num_rels * 2)
Mufei Li's avatar
Mufei Li committed
122
        self.reg_param = reg_param
123
        self.w_relation = nn.Parameter(torch.Tensor(num_rels, h_dim))
Mufei Li's avatar
Mufei Li committed
124
125
126
127
128
129
130
        nn.init.xavier_uniform_(self.w_relation,
                                gain=nn.init.calculate_gain('relu'))

    def calc_score(self, embedding, triplets):
        s = embedding[triplets[:,0]]
        r = self.w_relation[triplets[:,1]]
        o = embedding[triplets[:,2]]
131
        score = torch.sum(s * r * o, dim=1)
Mufei Li's avatar
Mufei Li committed
132
133
        return score

134
    def forward(self, g, nids):
135
        return self.rgcn(g, nids)
Mufei Li's avatar
Mufei Li committed
136
137

    def regularization_loss(self, embedding):
138
        return torch.mean(embedding.pow(2)) + torch.mean(self.w_relation.pow(2))
Mufei Li's avatar
Mufei Li committed
139
140
141
142
143
144
145
146

    def get_loss(self, embed, triplets, labels):
        # each row in the triplets is a 3-tuple of (source, relation, destination)
        score = self.calc_score(embed, triplets)
        predict_loss = F.binary_cross_entropy_with_logits(score, labels)
        reg_loss = self.regularization_loss(embed)
        return predict_loss + self.reg_param * reg_loss

147
148
149
150
151
152
153
154
155
156
157
158
159
160
def filter(triplets_to_filter, target_s, target_r, target_o, num_nodes, filter_o=True):
    """Get candidate heads or tails to score"""
    target_s, target_r, target_o = int(target_s), int(target_r), int(target_o)
    # Add the ground truth node first
    if filter_o:
        candidate_nodes = [target_o]
    else:
        candidate_nodes = [target_s]
    for e in range(num_nodes):
        triplet = (target_s, target_r, e) if filter_o else (e, target_r, target_o)
        # Do not consider a node if it leads to a real triplet
        if triplet not in triplets_to_filter:
            candidate_nodes.append(e)
    return torch.LongTensor(candidate_nodes)
Mufei Li's avatar
Mufei Li committed
161

162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def perturb_and_get_filtered_rank(emb, w, s, r, o, test_size, triplets_to_filter, filter_o=True):
    """Perturb subject or object in the triplets"""
    num_nodes = emb.shape[0]
    ranks = []
    for idx in tqdm.tqdm(range(test_size), desc="Evaluate"):
        target_s = s[idx]
        target_r = r[idx]
        target_o = o[idx]
        candidate_nodes = filter(triplets_to_filter, target_s, target_r,
                                 target_o, num_nodes, filter_o=filter_o)
        if filter_o:
            emb_s = emb[target_s]
            emb_o = emb[candidate_nodes]
        else:
            emb_s = emb[candidate_nodes]
            emb_o = emb[target_o]
        target_idx = 0
        emb_r = w[target_r]
        emb_triplet = emb_s * emb_r * emb_o
        scores = torch.sigmoid(torch.sum(emb_triplet, dim=1))
Mufei Li's avatar
Mufei Li committed
182

183
184
185
186
        _, indices = torch.sort(scores, descending=True)
        rank = int((indices == target_idx).nonzero())
        ranks.append(rank)
    return torch.LongTensor(ranks)
Mufei Li's avatar
Mufei Li committed
187

188
189
190
191
192
193
194
195
196
197
198
199
200
201
def calc_mrr(emb, w, test_mask, triplets_to_filter, batch_size=100, filter=True):
    with torch.no_grad():
        test_triplets = triplets_to_filter[test_mask]
        s, r, o = test_triplets[:,0], test_triplets[:,1], test_triplets[:,2]
        test_size = len(s)
        triplets_to_filter = {tuple(triplet) for triplet in triplets_to_filter.tolist()}
        ranks_s = perturb_and_get_filtered_rank(emb, w, s, r, o, test_size,
                                                triplets_to_filter, filter_o=False)
        ranks_o = perturb_and_get_filtered_rank(emb, w, s, r, o,
                                                test_size, triplets_to_filter)
        ranks = torch.cat([ranks_s, ranks_o])
        ranks += 1 # change to 1-indexed
        mrr = torch.mean(1.0 / ranks.float()).item()
    return mrr
Mufei Li's avatar
Mufei Li committed
202

203
204
def train(dataloader, test_g, test_nids, test_mask, triplets, device, model_state_file, model):
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
Mufei Li's avatar
Mufei Li committed
205
    best_mrr = 0
206
    for epoch, batch_data in enumerate(dataloader): # single graph batch
Mufei Li's avatar
Mufei Li committed
207
        model.train()
208
        g, train_nids, edges, labels = batch_data
Mufei Li's avatar
Mufei Li committed
209
        g = g.to(device)
210
211
        train_nids = train_nids.to(device)
        edges = edges.to(device)
Mufei Li's avatar
Mufei Li committed
212
213
        labels = labels.to(device)

214
215
        embed = model(g, train_nids)
        loss = model.get_loss(embed, edges, labels)
Mufei Li's avatar
Mufei Li committed
216
217
218
219
220
221
222
223
224
        optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # clip gradients
        optimizer.step()
        print("Epoch {:04d} | Loss {:.4f} | Best MRR {:.4f}".format(epoch, loss.item(), best_mrr))
        if (epoch + 1) % 500 == 0:
            # perform validation on CPU because full graph is too large
            model = model.cpu()
            model.eval()
225
            embed = model(test_g, test_nids)
Mufei Li's avatar
Mufei Li committed
226
            mrr = calc_mrr(embed, model.w_relation, test_mask, triplets,
227
                           batch_size=500)
Mufei Li's avatar
Mufei Li committed
228
229
230
            # save best model
            if best_mrr < mrr:
                best_mrr = mrr
231
                torch.save({'state_dict': model.state_dict(), 'epoch': epoch}, model_state_file)
Mufei Li's avatar
Mufei Li committed
232
233
            model = model.to(device)

234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'Training with DGL built-in RGCN module')

    # load and preprocess dataset
    data = FB15k237Dataset(reverse=False)
    g = data[0]
    num_nodes = g.num_nodes()
    num_rels = data.num_rels
    train_g = get_subset_g(g, g.edata['train_mask'], num_rels)
    test_g = get_subset_g(g, g.edata['train_mask'], num_rels, bidirected=True)
    test_g.edata['norm'] = dgl.norm_by_dst(test_g).unsqueeze(-1)
    test_nids = torch.arange(0, num_nodes)
    test_mask = g.edata['test_mask']
    subg_iter = SubgraphIterator(train_g, num_rels) # uniform edge sampling
    dataloader = GraphDataLoader(subg_iter, batch_size=1, collate_fn=lambda x: x[0])

    # Prepare data for metric computation
    src, dst = g.edges()
    triplets = torch.stack([src, g.edata['etype'], dst], dim=1)

    # create RGCN model
    model = LinkPredict(num_nodes, num_rels).to(device)

    # train
    model_state_file = 'model_state.pth'
    train(dataloader, test_g, test_nids, test_mask, triplets, device, model_state_file, model)

    # testing
    print("Testing...")
    checkpoint = torch.load(model_state_file)
Mufei Li's avatar
Mufei Li committed
265
266
267
    model = model.cpu() # test on CPU
    model.eval()
    model.load_state_dict(checkpoint['state_dict'])
268
    embed = model(test_g, test_nids)
269
270
271
    best_mrr = calc_mrr(embed, model.w_relation, test_mask, triplets,
                        batch_size=500)
    print("Best MRR {:.4f} achieved using the epoch {:04d}".format(best_mrr, checkpoint['epoch']))