Question about bipartite graphs

[zcaicaros]

Hi PyG creator and community,

I have a question regarding bipartite graph data in the following. It would be much appreciated if can get any hint from you!

By following this doc, I create a bipartite graph data like this:

from torch_geometric.data import Data
class BipartiteData(Data):
    def __init__(self, edge_index=None, x_s=None, x_t=None, edge_attr=None):
        super().__init__()
        self.edge_index = edge_index
        self.x = (x_s, x_t)
        self.edge_attr = edge_attr

Then I instantiate it like:

import torch
dev = 'cuda'
x_s = torch.rand(2, 1)
x_t = torch.rand(3, 1)
edge_index = torch.zeros(5, 5)
edge_index[0, 0] = 1
edge_index[0, 2] = 1
edge_index[1, 1] = 1
edge_index[1, 2] = 1
edge_index = edge_index.nonzero().t()
edge_x = torch.rand(edge_index.size(1), 1)

data = BipartiteData(edge_index=edge_index, x_s=x_s, x_t=x_t, edge_attr=edge_x).to(dev)

But when I access the num_nodes, it only counts in the total number of target nodes x_t with a warning:

>>>print("number of nodes is:{}".format(data.num_nodes))
UserWarning: Unable to accurately infer 'num_nodes' from the attribute set '{'edge_attr', 'x', 'edge_index'}'. Please explicitly set 'num_nodes' as an attribute of 'data' to suppress this warning
  warnings.warn(
number of nodes is:3

Q1: Is it an expected behavior to count only the number of target nodes?

To remove the warning, I add num_nodes as an attribute in the class class BipartiteData(Data) as suggested :

'''
self.num_nodes = x_s.size(0) + x_t.size(0)
'''

Now the warning is gone and the num_nodes counts in both source nodes x_s and target nodes x_t, which is 2+3=5:

number of nodes is: 5

So far so good, then a new problem happens when I try to create a batch using torch_geometric,data.Batch:

>>>from torch_geometric.data import Batch
>>>data_batch = Batch.from_data_list([data, data]).to(dev)
Error:
Traceback (most recent call last):
  File "C:\work\Constructive-DRL-CJSP\model\context_gat.py", line 343, in <module>
    data_batch = Batch.from_data_list([data, data]).to(dev)
  File "C:\work\env\lib\site-packages\torch_geometric\data\batch.py", line 69, in from_data_list
    batch, slice_dict, inc_dict = collate(
  File "C:\work\env\lib\site-packages\torch_geometric\data\collate.py", line 33, in collate
    out = cls(_base_cls=data_list[0].__class__)  # Dynamic inheritance.
  File "C:\work\env\lib\site-packages\torch_geometric\data\batch.py", line 41, in __call__
    return super(DynamicInheritance, new_cls).__call__(*args, **kwargs)
  File "C:\work\Constructive-DRL-CJSP\model\context_gat.py", line 271, in __init__
    self.num_nodes = x_s.size(0) + x_t.size(0)
AttributeError: 'NoneType' object has no attribute 'size'

Then I just remove the line self.num_nodes = x_s.size(0) + x_t.size(0) in class BipartiteData(Data), then the batch can be created with the aforementioned warning:

>>>print(data_batch)
>>>print("number of nodes is:{}".format(data_batch.num_nodes))
UserWarning: Unable to accurately infer 'num_nodes' from the attribute set '{'x', 'edge_index', 'edge_attr'}'. Please explicitly set 'num_nodes' as an attribute of 'data' to suppress this warning
  warnings.warn(
BipartiteDataBatch(x=[2], edge_index=[2, 8], edge_attr=[8, 1], batch=[6], ptr=[3])
number of nodes is:6

As expected, it only counts the target nodes which is 3+3=6.

Q2: So how to create bipartite graph data and its batch?

[rusty1s]

In general, data.num_nodes is undefined for bipartite graphs, and you should simply avoid to call it in this scenario. Instead, simply use x_s.size(0) and x_t.size(0) to infer the number of source and destination nodes. Alternatively, you can override num_nodes directly:

from torch_geometric.data import Data
class BipartiteData(Data):
    def __init__(self, edge_index=None, x_s=None, x_t=None, edge_attr=None):
        super().__init__()
        self.edge_index = edge_index
        self.x_s = x_s
        self.x_t = x_t
        self.edge_attr = edge_attr

    def __inc__(self, key, value, *args, **kwargs):
        if key == 'edge_index':
            return torch.tensor([[self.x_s.size(0)], [self.x_t.size(0)]])
        else:
            return super().__inc__(key, value, *args, **kwargs)
            
    @property
    def num_nodes(self):
        return self.x_s.size(0) + self.x_t.size(1)

[zcaicaros]

Thanks for your reply. Please bear with me for one more question, if I model the bipartite graph as its equivalent counterpart - an NxN graph where N=x_s.size(0) + x_t.size(0), will this cause any differences during training?

[rusty1s]

No that should work as well, but it requires that x_s and x_t share the same features.

[zcaicaros]

Thanks, but what do you mean by 'share the same features?'. I guess you mean the dimension of features of nodes in x_s and x_t should be equal such that in the NxN graph we can have x = torch.cat([x_s, x_t], dim=0)?

[rusty1s]

Yes, exactly.