Handling Edge Features in HeteroGraphConv wrapper for heterogeneous graphs

I am working with a Heterogeneous graph, and I am using the dgl.EGATConv module (a modified version to handle bipartite graphs) which basically modifies the node features as well as edge features at each layer. But I don’t think the ‘dgl HeteroGraphConv’ wrapper can handle the edge features right now. In the sense, it only takes a tuple of tensors (source and destination node features). My question is how can I pass the edge features to HeteroGraphConv. Thanks!

HeteroGraphConv allows you to pass arbitrary positional arguments and keyword arguments via mod_args and mod_kwargs of its forward function in the form of dictionaries with keys being the relation type and values being the extra arguments. In your case, you can provide a dictionary of edge feature tensors to it.

I understand thanks! Still, I would like to point you to the fact that EGATConv can not handle bipartite graphs and it returns three dictionaries - node_features, edge_features, attention_scores which the HeteroGraphConv can not handle. Isn’t it better that we create a separate hetero-wrapper for this and modify EGATConv also?
Another issue is, that the current implementation of EGATConv when applied with HeteroGraphConv on MFGs(blocks of edge loaders) the edge_features (as input) and the tensor from apply_edges() will have different shapes (axis 0).

@mufeili Can EGATConv handle bipartite graph in theory? If so, that may be something we could quickly fix.

it returns three dictionaries - node_features, edge_features, attention_scores which the HeteroGraphConv can not handle.

This is indeed a limitation of current HeteroGraphConv. I’m thinking of extending it. Welcome any suggestions on this.

Another issue is, that the current implementation of EGATConv when applied with HeteroGraphConv on MFGs(blocks of edge loaders) the edge_features (as input) and the tensor from apply_edges() will have different shapes (axis 0).

Could you post some code snippet?

EGATConv on bipartite graphs:
From a practical real-world implementation point of view, most models will be more valuable if they can be applied to Heterogeneous graphs (with multiple bi-partite meta path types) with edge features.

EGATConv with MFG(block) input:
If multiple blocks are created using an Edgeloader, based on the graph the num_dst_nodes, and num_edges will change(both will reduce even with full neighborhood sampler) from block[0] to block[n]. To incorporate this factor for node edges in HeteroGraphConv we have this logic-

src_inputs = inputs
dst_inputs = {k: v[:g.number_of_dst_nodes(k)] for k, v in inputs.items()}

but for edge features in EGATConv -

f_ni = self.fc_ni(nfeats)
f_nj = self.fc_nj(nfeats)
f_fij = self.fc_fij(efeats)
graph.srcdata.update({'f_ni': f_ni})
graph.dstdata.update({'f_nj': f_nj})
# add ni, nj factors
graph.apply_edges(fn.u_add_v('f_ni', 'f_nj', 'f_tmp'))
# add fij to node factor
f_out = graph.edata.pop('f_tmp') + f_fij

‘f_tmp’ will be applied to only the edges of the concerned MFG (graph = block) and will have dimensions (n, d1) but ‘f_ij’ (passed as input) will have the dimensions of (m,d2), where m is the number of all the edges of a particular edge type. This will break if n!=m, which is the case with blocks.
This can be solved
by doing this -
f_out = graph.edata.pop('f_tmp') + torch.index_select(f_fij, 0, graph.edata['_ID']
or using some numpy function, thanks to ‘_ID’

HeteroGraphConv Variations:
There could be only 4 (6) variations of the hetero wrapper based on the number of inputs and outputs. I think we should include these variations to cover all the possibilities.

input variations:

  1. graph, node_features
  2. graph, node_featues, edge_features

output variations:

  1. node_features
    2.a. node_features, edge_attentions
    2.b. node_features , edge_features, edge_attentions

I want to contribute to these problems and discuss the solutions. Please feel free to connect

@minjie In theory EGATConv can handle bipartite graphs.

@rijulizer It sounds like efeats passed as you described is the features of all edges in the original graph while it should only be the features of the edges in the MFG.

@BarclayII Do we have an example of dealing with edge features with MFGs sampled?

@mufeili you are right but there is no provision for that yet,

Even if you pass these inputs-

nfeats = block[0].srcdata['features']
efeats = block[0].edata['features']

to EGATConv with HeteroGraphConv you will face the same issue.
because this kind of filtration -

src_inputs = inputs
dst_inputs = {k: v[:g.number_of_dst_nodes(k)] for k, v in inputs.items()}

is not happening at HeteroGraphConv level and thus the solution I propose.

Edge features will appear in mfg.edata[key]. Is this what you are asking?

Edge features will appear in mfg.edata[key] . Is this what you are asking?

Yes. Thanks.

to EGATConv with HeteroGraphConv you will face the same issue.
because this kind of filtration -

src_inputs = inputs
dst_inputs = {k: v[:g.number_of_dst_nodes(k)] for k, v in inputs.items()}

is not happening at HeteroGraphConv level and thus the solution I propose.

You can do it before passing the input features to a HeteroGraphConv object. It seems that the only issue is the support of bipartite graphs in EGATConv.

Sorry I did not notice that EGATConv returns both updated node features and edge features. In that case, it’s hard to come up with a general solution at this moment. You will need to modify HeteroGraphConv yourself for now.

@mufeili @minjie Can we extend EGATConv to handle bipartite graphs and HeteroGraphConv to HeteroEGATConv?

We can extend EGATConv for handling bipartite graphs. It will be great if you can make a contribution for that. Meanwhile, I’m not sure if we can achieve a solution that is general enough for extending HeteroGraphConv to support a collection of NN modules beyond EGATConv.

1 Like