MetaLayer

[mdanb]

The docs have an example of how to use the MetaLayer as follows:

from torch.nn import Sequential as Seq, Linear as Lin, ReLU
from torch_scatter import scatter_mean
from torch_geometric.nn import MetaLayer

class NodeModel(torch.nn.Module):
    def __init__(self):
        super(NodeModel, self).__init__()
        self.node_mlp_1 = Seq(Lin(..., ...), ReLU(), Lin(..., ...))
        self.node_mlp_2 = Seq(Lin(..., ...), ReLU(), Lin(..., ...))

    def forward(self, x, edge_index, edge_attr, u, batch):
        # x: [N, F_x], where N is the number of nodes.
        # edge_index: [2, E] with max entry N - 1.
        # edge_attr: [E, F_e]
        # u: [B, F_u]
        # batch: [N] with max entry B - 1.
        row, col = edge_index
        out = torch.cat([x[row], edge_attr], dim=1)
        out = self.node_mlp_1(out)
        out = scatter_mean(out, col, dim=0, dim_size=x.size(0))
        out = torch.cat([x, out, u[batch]], dim=1)
        return self.node_mlp_2(out)
op = MetaLayer(NodeModel())
x, _, _ = op(x, edge_index, edge_attr, u, batch)

I omitted the GlobalModel() and EdgeModel() because I'm not using them for now

I tried running this on my data, but I get the error:

  File "/home/mdanb/models/utils.py", line 45, in cross_validate
    x, _, _ = model(batch.x.float(), batch.edge_index, batch.edge_attr.float(), batch.u, batch.batch)
  File "/home/mdanb/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/mdanb/.local/lib/python3.8/site-packages/torch_geometric/nn/meta.py", line 112, in forward
    x = self.node_model(x, edge_index, edge_attr, u, batch)
  File "/home/mdanb/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/mdanb/models/gnn.py", line 92, in forward
    out = torch.cat([x, out, u[batch]], dim=1)
TypeError: only integer tensors of a single element can be converted to an index

The problem seems to be coming from the u[batch]. Is this trying to replicate the global feature for the graph N times (where N is the number of nodes in a single graph) or what is it meaning to do exactly?

[rusty1s]

Exactly, u[batch] will replicate the global features for each graph N times. Since this seems to fail for you, how does batch look like in your case?

[mdanb]

This is from my debugger:
batch = {Batch: 6} Batch(batch=[22113], edge_attr=[4688752,1], edge_index=[2,4688752], u=[1], x=[22113, 1], y=[1])
Here, my batch size is 1. The batch attribute in this case is just a 22113 tensor of 0s.

[rusty1s]

If u is a torch.Tensor and batch is a torch.LongTensor only holding zeros, then u[batch] should work just fine.