to_hetero_with_bases

[Xurorx]

Hello pyg-team!

I have a question regarding to to_hetero_with_bases function. I have a HeteroData object look like this:

HeteroData(
contact={
Id=[2769],
Emb=[2769, 128],
num_nodes=2769
},
file={
Id=[5729],
Emb=[5729, 100],
num_nodes=5729
},
(contact, relation1, contact)={ edge_index=[2, 52341] },
(contact, relation2, file)={ edge_index=[2, 48534] },
(contact, relation3, file)={ edge_index=[2, 38044] },
(file, relation4, contact)={ edge_index=[2, 24436] },
(file, relation5, contact)={ edge_index=[2, 1148] },
......... etc.. (many relations)
)

class SAGE(torch.nn.Module):
    """
    Graph Sage model
    """
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout):
        super(SAGE, self).__init__()

        self.convs = torch.nn.ModuleList()
        if num_layers == 1:
            self.convs.append(
                SAGEConv(in_channels, out_channels, normalize=False))
        else:
            self.convs.append(SAGEConv(in_channels, hidden_channels))
            for _ in range(num_layers - 2):
                self.convs.append(SAGEConv(hidden_channels, hidden_channels))
            self.convs.append(SAGEConv(hidden_channels, out_channels))

        self.dropout = dropout

    def forward(self, emb, adj_t):
        for conv in self.convs[:-1]:
            emb = conv(emb, adj_t)
            emb = F.relu(emb)
            emb = F.dropout(emb, p=self.dropout, training=self.training)
        emb = self.convs[-1](emb, adj_t)
        return emb

class Model(torch.nn.Module):
    """
    Heterogeneous GNN model
    """
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers, dropout):
        super().__init__()

        # Instantiate homogeneous GNN:
        self.gnn = SAGE(in_channels=in_channels, hidden_channels=hidden_channels, out_channels=out_channels, num_layers=num_layers, dropout=dropout)

        # Convert GNN model into a heterogeneous variant:
        self.gnn = to_hetero_with_bases(self.gnn, metadata=data.metadata, num_bases=3)

For my data, node feature dim for contact is 128, while node feature dim for file is 100. I want to use to_hetero_with_bases for my SAGE model, may I ask how to configure in_channels for to_hetero_bases to let it accept different node feature dim?

self.gnn = to_hetero_with_bases(self.gnn, metadata=data.metadata, num_bases=3, in_channels=?)

[rusty1s]

Please take a look at https://github.com/pyg-team/pytorch_geometric/blob/master/test/nn/test_to_hetero_with_bases_transformer.py#L133.

[rusty1s]

The in_channels refer to the feature dimensionality you want emb to map to (to align with the in_channels of the first GNN layer for example). to_hetero_with_bases can well handle different feature dimensionalities across different node types. Here is another example.

[Xurorx]

Thank you @rusty1s for another example. But in this example, "model = to_hetero_with_bases(model, metadata, num_bases=3, in_channels={'x': 16})" I believe both "paper" and "author" feature has dim=16 and both their feature called x?

My case is different, where data['contact'].emb dim =128, data['file'].emb dim = 100, how should I set in_channels for to_hetero_with_bases? Should I set in_channels = {'data['contact'].emb': 128, 'data['file'].emb': 100}? That obviously looks incorrect. I think the key in in_channels need to be strings, I don't know how to give the keys in in_channels dict.

[Xurorx]

I have tried to set SAGEConv's in_channels = (-1,-1) for lazy initialization, and that works if I use to_hetero to convert model to heterogenous architecture.

But with SAGEConv's in_channels = (-1,-1), it does not work with self.gnn = to_hetero_with_bases(self.gnn, metadata=data.metadata, num_bases=3). It will throw errors: RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 128 but got size 100 for tensor number 1 in the list. Now I hope to_hetero_with_bases in_channels can handle that feature dim difference.

[rusty1s]

Yes, it can already handle this. For example, the linked test case already operates on node features of different size, see https://github.com/pyg-team/pytorch_geometric/blob/master/test/nn/test_to_hetero_with_bases_transformer.py#L117. The in_channels attribute in to_hetero_with_bases is just used to specify the shared embedding dim you want different node types to map to (that should be the dimension your GNN expects as input).

[Xurorx]

Thank you @rusty1s , I think I understand now. in_channels attribute in to_hetero_with_bases should specify the dim that we want to mapped to. And the key value in in_channels should be aligned with the forward function's input argument of homogeneous model we created. In my case, I should set in_channels= {'emb': 100} because SAGE class's forward function input is 'emb' and 'adj_t'.

[Xurorx]

Also I found a bug in this test file https://github.com/pyg-team/pytorch_geometric/blob/master/test/nn/test_to_hetero_with_bases_transformer.py#L117. If you take Net4() related code out and only test Net4() for to_hetero_with_bases, it will throw error. But if you test Net4() along with other models, the error disappeared shockingly! Also, when I test Net4() with below code and change in_channels = {'x': 16} to in_channels = {'x0': 16} the error disappeared. I don't quite understand the magic here.

from typing import Tuple

import torch
from torch import Tensor
from torch.nn import Linear, ReLU, Sequential

import torch_geometric.typing
from torch_geometric.nn import (
    GINEConv,
    MessagePassing,
    RGCNConv,
    SAGEConv,
    to_hetero_with_bases,
)
from torch_geometric.typing import SparseTensor


class Net4(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = SAGEConv(16, 16)
        self.conv2 = SAGEConv(16, 16)
        self.lin1 = Linear(3 * 16, 32)

    def forward(self, x0: Tensor, edge_index: Tensor) -> Tensor:
        x1 = self.conv1(x0, edge_index).relu_()
        x2 = self.conv2(x1, edge_index).relu_()
        return self.lin1(torch.cat([x0, x1, x2], dim=-1))


def test_to_hetero_with_bases():
    metadata = (['paper', 'author'], [('paper', 'cites', 'paper'),
                                      ('paper', 'written_by', 'author'),
                                      ('author', 'writes', 'paper')])

    x_dict = {'paper': torch.randn(100, 16), 'author': torch.randn(100, 8)}
    edge_index_dict = {
        ('paper', 'cites', 'paper'):
        torch.randint(100, (2, 200), dtype=torch.long),
        ('paper', 'written_by', 'author'):
        torch.randint(100, (2, 200), dtype=torch.long),
        ('author', 'writes', 'paper'):
        torch.randint(100, (2, 200), dtype=torch.long),
    }
    edge_attr_dict = {
        ('paper', 'cites', 'paper'): torch.randn(200, 8),
        ('paper', 'written_by', 'author'): torch.randn(200, 8),
        ('author', 'writes', 'paper'): torch.randn(200, 8),
    }

    model = Net4()
    in_channels = {'x': 16}
    model = to_hetero_with_bases(model, metadata, num_bases=4,
                                 in_channels=in_channels, debug=False)
    out = model(x_dict, edge_index_dict)
    assert isinstance(out, dict) and len(out) == 2
    assert out['paper'].size() == (100, 32)
    assert out['author'].size() == (100, 32)


test_to_hetero_with_bases()

[rusty1s]

That's an amazing finding. There was indeed a bug in to_hetero_with_bases, which modified the input in-place. As such, running the models sequentially only tested on node features with the same shared feature dimensionality.

Anyway, the bug is that one needs to use in_channels = {'x0': 16} rather than in_channels = {'x': 16} because the input argument is named x0. Fixed in #7363.

[Xurorx]

Nice! Thank you @rusty1s for always being responsive!