Link Prediction with use-edge-features set true returns error

Question

Hi!

I am using NeptuneML (setup using AWS provided Neptune ML cloud formation template). I am trying to train a link prediction model that does incorporate edge features such as weights for predictions but retrieving an error. In the model-hpo-configuration file I only change this parameter to true in fixed parameters.

If I don't change this parameter the training runs fine without errors but my the model predictions are not optimal. The edge features have properties such as weights that would help making predictions. How can I train my link prediction model while using edge features?

Here are the error logs

```
/opt/conda/lib/python3.6/site-packages/sklearn/base.py:315: UserWarning: Trying to unpickle estimator SimpleImputer from version 1.0.2 when using version 0.24.2. This might lead to breaking code or invalid results. Use at your own risk.
  UserWarning)

/opt/conda/lib/python3.6/site-packages/sklearn/base.py:315: UserWarning: Trying to unpickle estimator SimpleImputer from version 1.0.2 when using version 0.24.2. This might lead to breaking code or invalid results. Use at your own risk.
  UserWarning)

/codebuild/output/src740861454/src/aten/src/ATen/native/cuda/Indexing.cu:658: indexSelectLargeIndex: block: [125,0,0], thread: [96,0,0] Assertion `srcIndex < srcSelectDimSize` failed.

Traceback (most recent call last):
  File "run.py", line 95, in 
    trainer.run()
  File "/opt/conda/lib/python3.6/site-packages/graphlytics/trainer.py", line 198, in run
    trainer.fit()
  File "/opt/conda/lib/python3.6/site-packages/graphlytics/model_zoo/rgcn/rgcn_link_prediction.py", line 735, in fit
    self._train(args)
  File "/opt/conda/lib/python3.6/site-packages/graphlytics/model_zoo/rgcn/rgcn_link_prediction.py", line 932, in _train
    (train_seeds, valid_seeds))
  File "/opt/conda/lib/python3.6/site-packages/graphlytics/model_zoo/rgcn/utils.py", line 99, in decorated_function
    raise exception.__class__(trace)

RuntimeError: Traceback (most recent call last):
  File "/opt/conda/lib/python3.6/site-packages/graphlytics/model_zoo/rgcn/utils.py", line 87, in _queue_result
    res = func(*args, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/graphlytics/model_zoo/rgcn/rgcn_link_prediction.py", line 1096, in run
    feats = model(blocks, feats)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 756, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/graphlytics/model_zoo/rgcn/rgcn_link_prediction.py", line 156, in forward
    h = layer(block, h, block.edata['etype'], block.edata['norm'], efeat, efeat_bool)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 756, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/graphlytics/model_zoo/rgcn/model.py", line 1312, in forward
    fn.sum(msg='msg', out='h'))
  File "/opt/conda/lib/python3.6/site-packages/dgl/heterograph.py", line 4849, in update_all
    ndata = core.message_passing(g, message_func, reduce_func, apply_node_func)
  File "/opt/conda/lib/python3.6/site-packages/dgl/core.py", line 330, in message_passing
    msgdata = invoke_edge_udf(g, ALL, g.canonical_etypes[0], mfunc, orig_eid=orig_eid)
  File "/opt/conda/lib/python3.6/site-packages/dgl/core.py", line 85, in invoke_edge_udf
    return func(ebatch)
  File "/opt/conda/lib/python3.6/site-packages/graphlytics/model_zoo/rgcn/model.py", line 1168, in basis_message_func
    edge_feat_seeds = th.nonzero(edges.data["efeat_bool"]).squeeze()

RuntimeError: CUDA error: device-side assert triggered

```

Will appreciate any help!

re:Post Agent · Answer

The error you're encountering when trying to use edge features in your link prediction model with Neptune ML is likely due to a compatibility issue or a limitation in the current implementation.

Link prediction in Neptune ML is designed to predict the most likely destination nodes for a particular source node and outgoing edge, or the most likely source nodes for a given destination node and incoming edge. By default, it uses the graph structure and node embeddings to make these predictions.

While Neptune ML supports various types of features for nodes, including numerical, categorical, and text features, the current implementation may not fully support the use of edge features in link prediction tasks. This could explain why you're encountering an error when setting 'use-edge-features' to true.

To improve your link prediction model without relying on edge features, you could try the following approaches:

1. Ensure that you're using the most relevant node features in your model. Node properties can often capture information that might be represented in edge weights.

2. Experiment with different GNN architectures or hyperparameters. Neptune ML uses Relational Graph Convolutional Networks (R-GCNs) by default, but you might be able to improve performance by adjusting the number of layers, hidden dimensions, or other hyperparameters.

3. If your edge weights represent strength of connection, you could potentially incorporate this information by creating additional node features that aggregate information from connected edges.

4. Consider using a custom model implementation if the built-in models don't meet your specific requirements. Neptune ML allows users to define and implement custom models for specific scenarios.

If incorporating edge features is crucial for your use case, you may need to consult with AWS support or consider alternative graph machine learning frameworks that explicitly support edge features in link prediction tasks. As Neptune ML continues to evolve, support for edge features in link prediction may be added in future updates.  
**Sources**  
[Amazon Neptune ML for machine learning on graphs - Amazon Neptune](https://docs.aws.amazon.com/neptune/latest/userguide/machine-learning.html)  
[Models and model training in Amazon Neptune ML - Amazon Neptune](https://docs.aws.amazon.com/neptune/latest/userguide/machine-learning-models-and-training.html)

Update: Please open a Support case with AWS for Neptune team to provide insights into this.
Thanks!

Link Prediction with use-edge-features set true returns error

Relevant content