Sequential: graph pooling heterogeneous model

[wgeul]

I'd like to apply a graph pooling layer to a heterogeneous Sequential model. The PyTorch Geometric Sequential class provides an example for applying such a graph pooling layer (below), but I run into an error when I convert the model to a heterogeneous model with the torch_geometric.nn.to_hetero transformer.

Graph pooling:

from torch.nn import Linear, ReLU, Dropout
from torch_geometric.nn import Sequential, GCNConv, JumpingKnowledge
from torch_geometric.nn import global_mean_pool

model = Sequential('x, edge_index, batch', [
    (Dropout(p=0.5), 'x -> x'),
    (GCNConv(dataset.num_features, 64), 'x, edge_index -> x1'),
    ReLU(inplace=True),
    (GCNConv(64, 64), 'x1, edge_index -> x2'),
    ReLU(inplace=True),
    (lambda x1, x2: [x1, x2], 'x1, x2 -> xs'),
    (JumpingKnowledge("cat", 64, num_layers=2), 'xs -> x'),
    (global_mean_pool, 'x, batch -> x'),
    Linear(2 * 64, dataset.num_classes),
])

Transform to_hetero:

data = dataset[0]
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = to_hetero(model, data.metadata(), aggr='sum').to(device)

Where data consists of a HeteroData object with multiple node & edge types, e.g.:

HeteroData(
  y=[647],
  a={ x=[13805, 393] },
  b={ x=[375, 45] },
  (a, has, b)={ edge_index=[2, 13778] }
  (b, rev_has, a)={ edge_index=[2, 13778] },
)

Error:

  File "/tmp/ipykernel_3779/3918792047.py", line 1, in <module>
    to_hetero(model, data.metadata(), aggr='sum').to(device)

  File "lib/python3.7/site-packages/torch_geometric/nn/to_hetero_transformer.py", line 110, in to_hetero
    transformer = ToHeteroTransformer(module, metadata, aggr, input_map, debug)

  File "lib/python3.7/site-packages/torch_geometric/nn/to_hetero_transformer.py", line 134, in __init__
    super().__init__(module, input_map, debug)

  File "lib/python3.7/site-packages/torch_geometric/nn/fx.py", line 74, in __init__
    self.gm = symbolic_trace(module)

  File "lib/python3.7/site-packages/torch_geometric/nn/fx.py", line 245, in symbolic_trace
    return GraphModule(module, Tracer().trace(module, concrete_args))

  File "lib/python3.7/site-packages/torch/fx/_symbolic_trace.py", line 615, in trace
    self.create_node('output', 'output', (self.create_arg(fn(*args)),), {},

  File "tmp3btpwi63.py", line 24, in forward
    x = self.module_7(x, batch)

  File "lib/python3.7/site-packages/torch_geometric/nn/glob/glob.py", line 48, in global_mean_pool
    size = int(batch.max().item() + 1) if size is None else size

TypeError: int() argument must be a string, a bytes-like object or a number, not 'Proxy'

I'm assuming the error has to be related to the batch argument, which in the case of a heterogeneous model is a dictionary consisting of the node types as keys and the node type's batch tensor as items. In a homogeneous graph the batch object will be just a tensor.

Is it at all possible to apply graph pooling to a heterogeneous model? If so, how do I do this?

[rusty1s]

This is interesting. Thanks for reporting. Will look into supporting this soon :)

[rusty1s]

This is now integrated in master. See here for an example. Importantly, you have to use the newly introduced GlobalPooling module for this.

[wgeul]

Hi Matthias, thanks a lot for this swift reply and solution, very much appreciated! noted on the use of the GlobalPooling module.
Willem