nn.Sequential for heterogeneous graphs

[pstrybol]

Hi,

I wanted to dynamically build GNN models for heterogeneous graphs which is why I am using the Sequential module. Then, I have flexibility in testing the number and type of layers to benchmark the performance with. However, I also have a heterogeneous graph and it seems to conflict when using the Sequential module to build my models. Here is an example of 2 models, one build without using Sequential:

  (nt1_emb): Embedding(13953, 64)
  (nt2_emb): Embedding(35, 64)
  (gnn): GraphModule(
    (conv1): ModuleDict(
      (gene__interacts_with__gene): SAGEConv(64, 64, aggr=mean)
      (gene__dependency_of__cell): SAGEConv(64, 64, aggr=mean)
      (gene__rev_interacts_with__gene): SAGEConv(64, 64, aggr=mean)
      (cell__rev_dependency_of__gene): SAGEConv(64, 64, aggr=mean)
    )
    (conv2): ModuleDict(
      (gene__interacts_with__gene): SAGEConv(64, 64, aggr=mean)
      (gene__dependency_of__cell): SAGEConv(64, 64, aggr=mean)
      (gene__rev_interacts_with__gene): SAGEConv(64, 64, aggr=mean)
      (cell__rev_dependency_of__gene): SAGEConv(64, 64, aggr=mean)
    )
  )
  (classifier): LPsimple_classif()
)

And the same model architecture but now build using the Sequential module:

  (nt1_emb): Embedding(13953, 64)
  (nt2_emb): Embedding(35, 64)
  (gnn): GraphModule(
    (gnn_layers): ModuleDict(
      (gene__interacts_with__gene): Sequential(
        (0): SAGEConv(64, 64, aggr=mean)
        (1): ReLU()
        (2): Dropout(p=0.2, inplace=False)
        (3): SAGEConv(64, 64, aggr=mean)
        (4): ReLU()
        (5): Dropout(p=0.2, inplace=False)
      )
      (gene__dependency_of__cell): Sequential(
        (0): SAGEConv(64, 64, aggr=mean)
        (1): ReLU()
        (2): Dropout(p=0.2, inplace=False)
        (3): SAGEConv(64, 64, aggr=mean)
        (4): ReLU()
        (5): Dropout(p=0.2, inplace=False)
      )
      (gene__rev_interacts_with__gene): Sequential(
        (0): SAGEConv(64, 64, aggr=mean)
        (1): ReLU()
        (2): Dropout(p=0.2, inplace=False)
        (3): SAGEConv(64, 64, aggr=mean)
        (4): ReLU()
        (5): Dropout(p=0.2, inplace=False)
      )
      (cell__rev_dependency_of__gene): Sequential(
        (0): SAGEConv(64, 64, aggr=mean)
        (1): ReLU()
        (2): Dropout(p=0.2, inplace=False)
        (3): SAGEConv(64, 64, aggr=mean)
        (4): ReLU()
        (5): Dropout(p=0.2, inplace=False)
      )
    )
  )
  (classifier): LPsimple_classif()
)

I can see one subtle difference: in the first, a single ModuleDict contains all interaction types for a certain conv layer. In the second example, each ModuleDict contains all conv layers for a single edge type. If this is the problem do you have a suggestion of how to fix this?

Thanks!

[rusty1s]

Do you have a simple example to reproduce this? Are you just calling to_hetero on the final Sequential model?

[pstrybol]

Hi, sorry for the late reply.

I use a list comprehension for the generation of the list of GNN layers, then, I put them in an OrderedDict such that I can pass it to geom_nn.Sequential.

layers = [gat_layer(in_features=features[i], hidden_features=features[i+1],
                            act_fn=act_fn, dropout=self.dropout, heads=None, ix=i,
                            layer_name=layer_name, layer=gnn_factory[layer_name])
                  for i in range(self.num_layers)]
        
            layers = OrderedDict(chain(*[i.items() for i in layers]))
            self.gnn_layers = geom_nn.Sequential('x, edge_index_dict', layers)

Subsequently, in the forward function I call the following:

self.gnn = GCNcustom(features=features,
                                 layer_name=layer_name, heads=heads,
                                 dropout=dropout, act_fn=act_fn)

self.gnn = geom_nn.to_hetero(self.gnn, metadata=heterodata.metadata())

Mainly I wanted to use this to compare sageconv, and different GAT layers. Since using Sequential did not work at that time I build a separate GATmodel but here I encountered another problem. I could not return the attention weights because at some point in the call stack a tuple is passed because of the heterogeneous nature of the graph while it expects an int. Could it be that GATlayers are not yet fully compatible with heterogeneous graphs? The models trains fine when I don't return the attention weights.

Thanks!

[rusty1s]

I added a small test to confirm that to_hetero can work with Sequential PyG models (which runs just fine), please take a look: https://github.com/pyg-team/pytorch_geometric/pull/7927/files

[pstrybol]

HI, thanks for the example, I got it to work now, one small question, why are the in_features a tuple of (-1, -1). The only thing I struggle with is returning the attention weights.

[rusty1s]

I don't think this is currently possible with to_hetero. You would need to implement the heterogeneous model on your own in this case (or leverage message passing hooks for this).