Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit a9d69c7

Browse files
author
Hongyu Ren
committed
update argument
1 parent 9236d9c commit a9d69c7

4 files changed

Lines changed: 34 additions & 34 deletions

File tree

‎examples/lsc/wikikg90m/dgl-ke-ogb-lsc/python/dglke/eval.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ def main():
193193
n_relations = dataset.n_relations
194194
ckpt_path = args.model_path
195195
model = load_model_from_checkpoint(args, n_entities, n_relations, ckpt_path, dataset.entity_feat.shape[1], dataset.relation_feat.shape[1])
196-
if args.train_mode in ['roberta', 'both']:
196+
if args.train_mode in ['roberta', 'concat']:
197197
model.entity_feat.emb = dataset.entity_feat
198198
model.relation_feat.emb = dataset.relation_feat
199199

‎examples/lsc/wikikg90m/dgl-ke-ogb-lsc/python/dglke/models/general_models.py‎

Lines changed: 29 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -233,10 +233,10 @@ def __init__(self, args, model_name, n_entities, n_relations, hidden_dim, gamma,
233233
self.loss_gen = LossGenerator(args, args.loss_genre, args.neg_adversarial_sampling,
234234
args.adversarial_temperature, args.pairwise)
235235

236-
if self.train_mode in ['emb', 'both']:
236+
if self.train_mode in ['shallow', 'concat']:
237237
self.entity_emb = ExternalEmbedding(args, n_entities, entity_dim,
238238
F.cpu() if args.mix_cpu_gpu else device)
239-
if self.train_mode in ['roberta', 'both']:
239+
if self.train_mode in ['roberta', 'concat']:
240240
assert ent_feat_dim != -1 and rel_feat_dim != -1
241241
self.entity_feat = ExternalEmbedding(args, n_entities, ent_feat_dim,
242242
F.cpu() if args.mix_cpu_gpu else device, is_feat=True)
@@ -246,8 +246,8 @@ def __init__(self, args, model_name, n_entities, n_relations, hidden_dim, gamma,
246246
else:
247247
rel_dim = relation_dim
248248

249-
self.use_mlp = self.train_mode in ['both', 'roberta']
250-
if self.train_mode == 'both':
249+
self.use_mlp = self.train_mode in ['concat', 'roberta']
250+
if self.train_mode == 'concat':
251251
self.transform_net = MLP(entity_dim+ent_feat_dim, entity_dim, relation_dim+rel_feat_dim, relation_dim)
252252
# self.transform_e_net = torch.nn.Linear(entity_dim, entity_dim)
253253
# self.transform_r_net = torch.nn.Linear(relation_dim, relation_dim)
@@ -261,10 +261,10 @@ def __init__(self, args, model_name, n_entities, n_relations, hidden_dim, gamma,
261261
print(self.strict_rel_part, self.soft_rel_part)
262262
assert not self.strict_rel_part and not self.soft_rel_part
263263
if not self.strict_rel_part and not self.soft_rel_part:
264-
if self.train_mode in ['emb', 'both']:
264+
if self.train_mode in ['shallow', 'concat']:
265265
self.relation_emb = ExternalEmbedding(args, n_relations, rel_dim,
266266
F.cpu() if args.mix_cpu_gpu else device)
267-
if self.train_mode in ['roberta', 'both']:
267+
if self.train_mode in ['roberta', 'concat']:
268268
self.relation_feat = ExternalEmbedding(args, n_relations, rel_feat_dim,
269269
F.cpu() if args.mix_cpu_gpu else device, is_feat=True)
270270
else:
@@ -303,16 +303,16 @@ def __init__(self, args, model_name, n_entities, n_relations, hidden_dim, gamma,
303303
def share_memory(self):
304304
"""Use torch.tensor.share_memory_() to allow cross process embeddings access.
305305
"""
306-
if self.train_mode in ['both', 'emb']:
306+
if self.train_mode in ['concat', 'shallow']:
307307
self.entity_emb.share_memory()
308-
if self.train_mode in ['both', 'roberta']:
308+
if self.train_mode in ['concat', 'roberta']:
309309
self.entity_feat.share_memory()
310310
if self.strict_rel_part or self.soft_rel_part:
311311
self.global_relation_emb.share_memory()
312312
else:
313-
if self.train_mode in ['both', 'emb']:
313+
if self.train_mode in ['concat', 'shallow']:
314314
self.relation_emb.share_memory()
315-
if self.train_mode in ['both', 'roberta']:
315+
if self.train_mode in ['concat', 'roberta']:
316316
self.relation_feat.share_memory()
317317

318318
if self.model_name == 'TransR':
@@ -331,14 +331,14 @@ def save_emb(self, path, dataset):
331331
dataset : str
332332
Dataset name as prefix to the saved embeddings.
333333
"""
334-
if self.train_mode in ['emb', 'both']:
334+
if self.train_mode in ['shallow', 'concat']:
335335
self.entity_emb.save(path, dataset+'_'+self.model_name+'_entity')
336-
if self.train_mode in ['roberta', 'both']:
336+
if self.train_mode in ['roberta', 'concat']:
337337
torch.save({'transform_state_dict': self.transform_net.state_dict()}, os.path.join(path, dataset+"_"+self.model_name+"_mlp"))
338338
if self.strict_rel_part or self.soft_rel_part:
339339
self.global_relation_emb.save(path, dataset+'_'+self.model_name+'_relation')
340340
else:
341-
if self.train_mode in ['emb', 'both']:
341+
if self.train_mode in ['shallow', 'concat']:
342342
self.relation_emb.save(path, dataset+'_'+self.model_name+'_relation')
343343

344344
self.score_func.save(path, dataset+'_'+self.model_name)
@@ -360,11 +360,11 @@ def load_emb(self, path, dataset):
360360
def reset_parameters(self):
361361
"""Re-initialize the model.
362362
"""
363-
if self.train_mode in ['emb', 'both']:
363+
if self.train_mode in ['shallow', 'concat']:
364364
self.entity_emb.init(self.emb_init)
365365
self.score_func.reset_parameters()
366366
if (not self.strict_rel_part) and (not self.soft_rel_part):
367-
if self.train_mode in ['emb', 'both']:
367+
if self.train_mode in ['shallow', 'concat']:
368368
self.relation_emb.init(self.emb_init)
369369
else:
370370
self.global_relation_emb.init(self.emb_init)
@@ -424,9 +424,9 @@ def predict_neg_score(self, pos_g, neg_g, to_device=None, gpu_id=-1, trace=False
424424
neg_head_ids = neg_g.ndata['id'][neg_g.head_nid]
425425
if self.train_mode == 'roberta':
426426
neg_head = self.transform_net.embed_entity(self.entity_feat(neg_head_ids, gpu_id, False))
427-
elif self.train_mode == 'emb':
427+
elif self.train_mode == 'shallow':
428428
neg_head = self.entity_emb(neg_head_ids, gpu_id, trace)
429-
elif self.train_mode == 'both':
429+
elif self.train_mode == 'concat':
430430
neg_head = self.transform_net.embed_entity(torch.cat([self.entity_feat(neg_head_ids, gpu_id, False), self.entity_emb(neg_head_ids, gpu_id, trace)], -1))
431431

432432
head_ids, tail_ids = pos_g.all_edges(order='eid')
@@ -456,9 +456,9 @@ def predict_neg_score(self, pos_g, neg_g, to_device=None, gpu_id=-1, trace=False
456456
neg_tail_ids = neg_g.ndata['id'][neg_g.tail_nid]
457457
if self.train_mode == 'roberta':
458458
neg_tail = self.transform_net.embed_entity(self.entity_feat(neg_tail_ids, gpu_id, False))
459-
elif self.train_mode == 'emb':
459+
elif self.train_mode == 'shallow':
460460
neg_tail = self.entity_emb(neg_tail_ids, gpu_id, trace)
461-
elif self.train_mode == 'both':
461+
elif self.train_mode == 'concat':
462462
neg_tail = self.transform_net.embed_entity(torch.cat([self.entity_feat(neg_tail_ids, gpu_id, False), self.entity_emb(neg_tail_ids, gpu_id, trace)], -1))
463463

464464
head_ids, tail_ids = pos_g.all_edges(order='eid')
@@ -514,11 +514,11 @@ def predict_score_wikikg(self, query, candidate, mode, to_device=None, gpu_id=-1
514514
neg_head = self.transform_net.embed_entity(self.entity_feat(candidate.view(-1), gpu_id, False))
515515
tail = self.transform_net.embed_entity(self.entity_feat(query[:,0], gpu_id, False))
516516
rel = self.transform_net.embed_relation(self.relation_feat(query[:,1], gpu_id, False))
517-
elif self.train_mode == 'emb':
517+
elif self.train_mode == 'shallow':
518518
neg_head = self.entity_emb(candidate.view(-1), gpu_id, False)
519519
tail = self.entity_emb(query[:,0], gpu_id, False)
520520
rel = self.relation_emb(query[:,1], gpu_id, False)
521-
elif self.train_mode == 'both':
521+
elif self.train_mode == 'concat':
522522
neg_head = self.transform_net.embed_entity(torch.cat([self.entity_feat(candidate.view(-1), gpu_id, False), self.entity_emb(candidate.view(-1), gpu_id, False)], -1))
523523
tail = self.transform_net.embed_entity(torch.cat([self.entity_feat(query[:,0], gpu_id, False), self.entity_emb(query[:,0], gpu_id, False)], -1))
524524
rel = self.transform_net.embed_relation(torch.cat([self.relation_feat(query[:,1], gpu_id, False), self.relation_emb(query[:,1], gpu_id, False)], -1))
@@ -530,11 +530,11 @@ def predict_score_wikikg(self, query, candidate, mode, to_device=None, gpu_id=-1
530530
neg_tail = self.transform_net.embed_entity(self.entity_feat(candidate.view(-1), gpu_id, False))
531531
head = self.transform_net.embed_entity(self.entity_feat(query[:,0], gpu_id, False))
532532
rel = self.transform_net.embed_relation(self.relation_feat(query[:,1], gpu_id, False))
533-
elif self.train_mode == 'emb':
533+
elif self.train_mode == 'shallow':
534534
neg_tail = self.entity_emb(candidate.view(-1), gpu_id, False)
535535
head = self.entity_emb(query[:,0], gpu_id, False)
536536
rel = self.relation_emb(query[:,1], gpu_id, False)
537-
elif self.train_mode == 'both':
537+
elif self.train_mode == 'concat':
538538
neg_tail = self.transform_net.embed_entity(torch.cat([self.entity_feat(candidate.view(-1), gpu_id, False), self.entity_emb(candidate.view(-1), gpu_id, False)], -1))
539539
head = self.transform_net.embed_entity(torch.cat([self.entity_feat(query[:,0], gpu_id, False), self.entity_emb(query[:,0], gpu_id, False)], -1))
540540
rel = self.transform_net.embed_relation(torch.cat([self.relation_feat(query[:,1], gpu_id, False), self.relation_emb(query[:,1], gpu_id, False)], -1))
@@ -568,10 +568,10 @@ def forward(self, pos_g, neg_g, gpu_id=-1):
568568
if self.train_mode == 'roberta':
569569
pos_g.ndata['emb'] = self.transform_net.embed_entity(self.entity_feat(pos_g.ndata['id'], gpu_id, False))
570570
pos_g.edata['emb'] = self.transform_net.embed_relation(self.relation_feat(pos_g.edata['id'], gpu_id, False))
571-
elif self.train_mode == 'emb':
571+
elif self.train_mode == 'shallow':
572572
pos_g.ndata['emb'] = self.entity_emb(pos_g.ndata['id'], gpu_id, True)
573573
pos_g.edata['emb'] = self.relation_emb(pos_g.edata['id'], gpu_id, True)
574-
elif self.train_mode == 'both':
574+
elif self.train_mode == 'concat':
575575
pos_g.ndata['emb'] = self.transform_net.embed_entity(torch.cat([self.entity_feat(pos_g.ndata['id'], gpu_id, False), self.entity_emb(pos_g.ndata['id'], gpu_id, True)], -1))
576576
pos_g.edata['emb'] = self.transform_net.embed_relation(torch.cat([self.relation_feat(pos_g.edata['id'], gpu_id, False), self.relation_emb(pos_g.edata['id'], gpu_id, True)], -1))
577577
self.score_func.prepare(pos_g, gpu_id, True)
@@ -596,7 +596,7 @@ def forward(self, pos_g, neg_g, gpu_id=-1):
596596
loss, log = self.loss_gen.get_total_loss(pos_score, neg_score, edge_weight)
597597
# regularization: TODO(zihao)
598598
#TODO: only reg ent&rel embeddings. other params to be added.
599-
if self.args.regularization_coef > 0.0 and self.args.regularization_norm > 0 and self.train_mode in ['both', 'emb']:
599+
if self.args.regularization_coef > 0.0 and self.args.regularization_norm > 0 and self.train_mode in ['concat', 'shallow']:
600600
coef, nm = self.args.regularization_coef, self.args.regularization_norm
601601
reg = coef * (norm(self.entity_emb.curr_emb(), nm) + norm(self.relation_emb.curr_emb(), nm))
602602
log['regularization'] = get_scalar(reg)
@@ -610,7 +610,7 @@ def update(self, gpu_id=-1):
610610
gpu_id : int
611611
Which gpu to accelerate the calculation. if -1 is provided, cpu is used.
612612
"""
613-
if self.train_mode in ['emb', 'both']:
613+
if self.train_mode in ['shallow', 'concat']:
614614
self.entity_emb.update(gpu_id)
615615
self.relation_emb.update(gpu_id)
616616
self.score_func.update(gpu_id)
@@ -668,13 +668,13 @@ def load_relation(self, device=None):
668668
def create_async_update(self):
669669
"""Set up the async update for entity embedding.
670670
"""
671-
if self.train_mode in ['emb', 'both']:
671+
if self.train_mode in ['shallow', 'concat']:
672672
self.entity_emb.create_async_update()
673673

674674
def finish_async_update(self):
675675
"""Terminate the async update for entity embedding.
676676
"""
677-
if self.train_mode in ['emb', 'both']:
677+
if self.train_mode in ['shallow', 'concat']:
678678
self.entity_emb.finish_async_update()
679679

680680

‎examples/lsc/wikikg90m/dgl-ke-ogb-lsc/python/dglke/train.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,8 @@ def __init__(self):
6868
'as pos_score = pos_score * edge_importance')
6969

7070
self.add_argument('--print_on_screen', action='store_true')
71-
self.add_argument('--train_mode', type=str, default='emb',
72-
help='emb or roberta or both')
71+
self.add_argument('--train_mode', type=str, default='shallow',
72+
help='shallow or roberta or concat')
7373
self.add_argument('--mlp_lr', type=float, default=0.0001,
7474
help='The learning rate of optimizing mlp')
7575
self.add_argument('--seed', type=int, default=0,
@@ -300,7 +300,7 @@ def main():
300300
print("To create model")
301301
t1 = time.time()
302302
model = load_model(args, dataset.n_entities, dataset.n_relations, dataset.entity_feat.shape[1], dataset.relation_feat.shape[1])
303-
if args.train_mode in ['roberta', 'both']:
303+
if args.train_mode in ['roberta', 'concat']:
304304
model.entity_feat.emb = dataset.entity_feat
305305
model.relation_feat.emb = dataset.relation_feat
306306
print("Model created, it takes %s seconds" % (time.time()-t1))

‎examples/lsc/wikikg90m/dgl-ke-ogb-lsc/python/dglke/train_pytorch.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ def train(args, model, train_sampler, valid_samplers=None, test_samplers=None, r
7979
if args.soft_rel_part:
8080
model.prepare_cross_rels(cross_rels)
8181

82-
if args.train_mode in ['roberta', 'both']:
82+
if args.train_mode in ['roberta', 'concat']:
8383
model.transform_net = model.transform_net.to(th.device('cuda:' + str(gpu_id)))
8484
optimizer = th.optim.Adam(model.transform_net.parameters(), args.mlp_lr)
8585
else:

0 commit comments

Comments
 (0)