a running error in the link prediction task based on TransE model

[NatsuHB]

🐛 Bug

I found an error when I try to test the TransE model. The relative code is follow:

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', '-m', default='TransE', type=str, help='name of models')
    parser.add_argument('--gpu', '-g', default='-1', type=int, help='-1 means cpu')
    parser.add_argument('--mini-batch-flag', action='store_true')

    args = parser.parse_args()

    ds = MyLPDatasetWithPredEdges()
    new_ds = AsLinkPredictionDataset(ds, target_link=target_link, target_link_r=target_link_r,
                                     split_ratio=[0.8, 0.1, 0.1], force_reload=True)

    experiment = Experiment(conf_path='../config.ini', max_epoch=3, model=args.model, dataset=new_ds,
                            task='link_prediction', mini_batch_flag=args.mini_batch_flag, gpu=args.gpu,
                            test_flag=True, prediction_flag=True, batch_size=100)
    experiment.run()

When the training process is over and the test process is about to begin, the following error occurs:

------------------------------------------------------------------------------
 Basic setup of this experiment: 
     model: TransE    
     dataset: my-lp-dataset-as-linkpred   
     task: link_prediction. 
 This experiment has following parameters. You can use set_params to edit them.
 Use print(experiment) to print this information again.
------------------------------------------------------------------------------
batch_size: 100
dataset_name: my-lp-dataset-as-linkpred
device: cpu
dis_norm: 1
filtered: filtered
gpu: -1
hidden_dim: 400
hpo_search_space: None
hpo_trials: 100
load_from_pretrained: False
lr: 1.0
margin: 4.0
max_epoch: 3
mini_batch_flag: False
model_name: TransE
neg_size: 13
optimizer: SGD
output_dir: ./openhgnn/output\TransE
patience: 3
prediction_flag: False
score_fn: transe
seed: 0
test_flag: False
test_percent: 0.1
use_best_config: False
valid_percent: 0.01
weight_decay: 0.0001

17 Feb 16:50    INFO  [Init Task] The task: link prediction, the dataset: Dataset("my-lp-dataset-as-linkpred", num_graphs=1, save_path=\.dgl\my-lp-dataset-as-linkpred), the evaluation metric is roc_auc, the score function: transe 
17 Feb 16:50    INFO  [Train Info] epoch 000
100%|███████████████████████████████████████████████████████████| 1011/1011 [01:00<00:00, 16.65it/s]
17 Feb 16:51    INFO  [Train Info] epoch 000 loss: 1749.946214914322
Traceback (most recent call last):
  File "./OpenHGNN/examples/applications/link_prediction/train.py", line 21, in <module>
    experiment.run()
  File ".\OpenHGNN\openhgnn\experiment.py", line 107, in run
    result = flow.train()
  File ".\OpenHGNN\openhgnn\trainerflow\TransX_trainer.py", line 53, in train
    epoch = self._train()
  File ".\OpenHGNN\openhgnn\trainerflow\TransX_trainer.py", line 82, in _train
    val_metric = self._test_step('valid')
  File ".\OpenHGNN\openhgnn\trainerflow\TransX_trainer.py", line 127, in _test_step
    return {mode: self.task.evaluate(n_emb, r_emb, mode)}
  File ".\OpenHGNN\openhgnn\tasks\link_prediction.py", line 129, in evaluate
    p_score = th.sigmoid(self.ScorePredictor(eval_hg, n_embedding, r_embedding))
  File ".\OpenHGNN\openhgnn\models\TransE.py", line 36, in forward
    h_emb = self.n_emb(h.to(self.device))
  File "\Anaconda3\envs\threatrace\lib\site-packages\torch\nn\modules\module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "\Anaconda3\envs\threatrace\lib\site-packages\torch\nn\modules\sparse.py", line 162, in forward
    self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "\Anaconda3\envs\threatrace\lib\site-packages\torch\nn\functional.py", line 2210, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
TypeError: embedding(): argument 'indices' (position 2) must be Tensor, not DGLGraph

exit code 1

Environment

• OpenHGNN Version =1.0:
• PyTorch 1.13.1+cu116, DGL 1.0.0:
• OS Windows:
• Python version: 3.7.16

[lazishu2000]

Use python main.py -m TransE -t link_prediction -d FB15k -g 0 --use_best_config and it should be fine. Thanks for comments.