How to input edge_attr into heterogeneous graph model

[RENEK-bool]

Here I excerpt an example from the docs. If I want to input edge_attr into the model, how do I modify the following code?

    def forward(self, x_dict, edge_index_dict):
        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)
            x_dict = {key: x.relu() for key, x in x_dict.items()}
        return self.lin(x_dict['author'])

How about modifying the code like this? But this method doesn’t seem to work. Because I tried to do this on my own dataset, but I got an error about dimension mismatch.

    def forward(self, x_dict, edge_index_dict, edge_attr_dict):
        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict, edge_attr_dict)
            x_dict = {key: x.relu() for key, x in x_dict.items()}
        return self.lin(x_dict['author'])

[rusty1s]

This generally works, but it depends a bit how your GNN layers are defined. Can you share more details about how your model __init__ is defined, and how your input data looks like?

[RENEK-bool]

Below is the heterogeneous graph I constructed

HeteroData(
  pin={ x=[2820, 3] },
  (pin, cell_arc, pin)={
    edge_index=[2, 1804],
    edge_attr=[1804, 1],
    y=[1804],
  },
  (pin, net_arc, pin)={
    edge_index=[2, 1623],
    edge_attr=[1623, 15],
    y1=[1623],
    y2=[1623],
  }
)

Then this is my model.

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        
        self.wire_convs = torch.nn.ModuleList()
        for _ in range(3):
            conv = HeteroConv({
                ('pin', 'cell_arc', 'pin'): GATConv((-1, -1), hidden_channels),
                ('pin', 'net_arc', 'pin'): GATConv((-1, -1), hidden_channels),
            }, aggr='sum')
            self.wire_convs.append(conv)

        self.act = nn.ReLU()
        self.wire_fc1 = nn.Linear(hidden_channels,32)
        self.wire_fc2 = nn.Linear(32,15)

    def forward(self, wire_dict, wire_edge_index_dict):
        
        for wire_conv in self.wire_convs:
            wire_dict = wire_conv(wire_dict, wire_edge_index_dict)
            wire_dict = {key: x.relu() for key, x in wire_dict.items()}
        
        x = wire_dict['pin']
        x = self.wire_fc1(x)
        x = self.act(x)
        x = self.wire_fc2(x)
    
        return x

Just now I discovered that an error will be reported when replacing GATConv with TransformerConv. If it is GATConv, the code can run. The following is the error message when TransformerConv.

Traceback (most recent call last):
  File "/home/gdut/Documents/ZJY/MTL_wirelength_timing/zjy8.py", line 125, in <module>
    out = model(wire_train_sample.x_dict, wire_train_sample.edge_index_dict, wire_train_sample.edge_attr_dict)
  File "/home/gdut/anaconda3/envs/zjy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/gdut/anaconda3/envs/zjy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/gdut/Documents/ZJY/MTL_wirelength_timing/zjy8.py", line 107, in forward
    wire_dict = wire_conv(wire_dict, wire_edge_index_dict, wire_edge_attr_dict)
  File "/home/gdut/anaconda3/envs/zjy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/gdut/anaconda3/envs/zjy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/gdut/anaconda3/envs/zjy/lib/python3.9/site-packages/torch_geometric/nn/conv/hetero_conv.py", line 158, in forward
    out = conv(*args, **kwargs)
  File "/home/gdut/anaconda3/envs/zjy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/gdut/anaconda3/envs/zjy/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/gdut/anaconda3/envs/zjy/lib/python3.9/site-packages/torch_geometric/nn/conv/transformer_conv.py", line 247, in forward
    out = out + x_r
RuntimeError: The size of tensor a (5087280) must match the size of tensor b (2820) at non-singleton dimension 0

[rusty1s]

You need to also specify edge_dim=-1 when defining your GNN layers.

[jykr]

Hi, I have a related question- during the batch sampling, my edge_attr is not subsampled together with edge_index, am I missing something here or should I manually subset with e_id?

Here's my data:

HeteroData(
  source={
    x=[2700, 50],
    size_factor=[2700],
    cont_cov=[2700],
    n_id=[2700],
  },
  target={
    size_factor=[1738],
    x=[1738, 50],
    n_id=[1738],
  },
  (source, rel, target)={
    edge_index=[2, 541461],
    edge_attr=[541461],
    edge_dist='Normal',
  },
  (target, rev_rel, source)={
    edge_index=[2, 541461],
    edge_attr=[541461],
  }
)

Loader:

NeighborLoader(
        train_data,
        num_neighbors={key: [30] * 2 for key in bidirectional_data.edge_types},
        batch_size=128,
        input_nodes=nodetype,
    )

Batch:

batch_bd HeteroData(
  source={
    x=[2699, 50],
    size_factor=[2699],
    cont_cov=[2699],
    n_id=[2699],
    num_sampled_nodes=[3],
    input_id=[128],
    batch_size=128,
  },
  target={
    size_factor=[1187],
    x=[1187, 50],
    n_id=[1187],
    num_sampled_nodes=[3],
  },
  (source, rel, target)={
    edge_index=[2, 35345],
    edge_dist='Normal',
    edge_attr=[146194],
    edge_attr_index=[2, 146194],
    e_id=[35345],
    num_sampled_edges=[2],
  },
  (target, rev_rel, source)={
    edge_index=[2, 3840],
    edge_attr=[146194],
    edge_attr_index=[2, 146194],
    e_id=[3840],
    num_sampled_edges=[2],
  }

[RENEK-bool]

@rusty1s

[rusty1s]

This is very strange, and I cannot reproduce this. Where is edge_attr_index coming from?