Is there a working example of Heterogeneous graph classification?

[carlosSorcero]

Hello all,

We have been using PyG for a while, mostly on homogeneous graphs for graph classification. We are now expanding the richness of our graph representations to include multiple node types and edge types. We still want to do graph classification but we are finding it hard to put a graph model together. Is there a working example of heterogeneous graph classification in PyG?

We took inspiration from the heterogeneous model in the documentation and the notebook for homogeneous graph classification.

Below is the code we've tried. However, we get the following error:

Traceback (most recent call last):
  File "graph_training.py", line 62, in <module>
    train(model, train_loader, optimizer, criterion, device)
  File "/home/paperspace/sorcero/gnn_disambiguation/hetero_training_example/training_utils.py", line 142, in train
    loss = criterion(out, data.y)  # Compute the loss.
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 603, in forward
    return F.binary_cross_entropy(input, target, weight=self.weight, reduction=self.reduction)
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch/nn/functional.py", line 2906, in binary_cross_entropy
    raise ValueError(
ValueError: Using a target size (torch.Size([8])) that is different to the input size (torch.Size([16, 1])) is deprecated. Please ensure they have the same size.

However, sometimes we get the following error:

Traceback (most recent call last):
  File "graph_training.py", line 61, in <module>
    train(model, train_loader, optimizer, criterion, device)
  File "/home/paperspace/sorcero/gnn_disambiguation/hetero_training_example/training_utils.py", line 140, in train
    out = predict(model, data.x_dict, data.edge_index_dict, data.batch_dict)
  File "/home/paperspace/sorcero/gnn_disambiguation/hetero_training_example/training_utils.py", line 132, in predict
    return model(x, edge_index, batch)
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/paperspace/sorcero/gnn_disambiguation/hetero_training_example/training_utils.py", line 116, in forward
    x_dict = conv(x_dict, edge_index_dict)
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch_geometric/nn/conv/hgt_conv.py", line 142, in forward
    k_dict[node_type] = self.k_lin[node_type](x).view(-1, H, D)
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch_geometric/nn/dense/linear.py", line 109, in forward
    return F.linear(x, self.weight, self.bias)
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch/nn/functional.py", line 1848, in linear
    return torch._C._nn.linear(input, weight, bias)
TypeError: linear(): argument 'input' (position 1) must be Tensor, not NoneType

Any help would really be appreciated -even a clean, simple working example of heterogeneous model classification.

HGT model and training functions

import torch
from torch_geometric.nn import global_mean_pool, HGTConv
from torch_geometric.nn import HANConv, GCNConv, SAGEConv
import torch.nn.functional as F
from torch.nn import Linear
from sklearn.metrics import confusion_matrix, precision_recall_fscore_support

class HAN(torch.nn.Module):
    def __init__(self, in_channels,
                 out_channels, hidden_channels=128, heads=8, metadata=None, num_classes=1):
        super().__init__()
        self.metadata = metadata
        self.hidden_channels = hidden_channels
        self.han_conv = HANConv(in_channels, self.hidden_channels, heads=heads,
                                dropout=dropout, metadata=metadata)
        self.lin1 = Linear(self.hidden_channels, out_channels)
        self.lin2 = Linear(out_channels, num_classes)

    def forward(self, x_dict, edge_index_dict, batch):
        out = self.han_conv(x_dict, edge_index_dict)

        concat = []
        for node_type in self.metadata[0]:
            concat.append(global_mean_pool(out[node_type], batch[node_type]))
        x_tensor = torch.reshape(torch.stack(concat), (-1, self.hidden_channels)) # this is wrong, we should keep the batch size shape
        x = F.dropout(x_tensor, p=dropout, training=self.training)
        x = self.lin1(x).relu()
        x = torch.sigmoid(self.lin2(x))

        return out

def predict(model, x, edge_index, batch):
    return model(x, edge_index, batch)


def train(model, loader, optimizer, criterion, device):
    model.train()

    for data in loader:  # Iterate in batches over the training dataset.
        data.to(device)
        out = predict(model, data.x_dict, data.edge_index_dict, data.batch_dict)
        # convert outout to tensor
        loss = criterion(out, data.y)  # Compute the loss.
        loss.backward()  # Derive gradients.
        optimizer.step()  # Update parameters based on gradients.
        optimizer.zero_grad()  # Clear gradients.

def test(model, loader, device):
    model.eval()

    losses = []
    y_trues = []
    y_preds = []
    for data in loader:  # Iterate in batches over the training/test dataset.
        data.to(device)
        out = predict(model, data.x_dict, data.edge_index_dict, data.batch_dict)
        y_true = data.y.cpu().numpy()
        y_pred = out.detach().cpu().numpy()
        loss = np.array([abs(o-e) for o, e in zip(y_pred, y_true)])
        losses.append(np.mean(loss))
        y_trues.extend(y_true)
        y_preds.extend(y_pred)
    return y_trues, y_preds, np.mean(losses)

Code to load up a fake dataset and create the model training cycle:

import random

import torch
from torch_geometric.data import HeteroData
from torch_geometric.loader import DataLoader
from sklearn.model_selection import train_test_split
from torch_geometric.datasets import FakeHeteroDataset
from torch_geometric.nn import to_hetero

from training_utils import HGT, HAN, GCN, predict, train, test, calculate_metrics


num_graphs = 10
avg_num_nodes = 20
batch_size = 16
device = 'cuda:0'
num_classes = 1
hidden_channels = 128
num_heads = 8
num_layers = 2
out_channels = 16


# Create test hetero graphs
graphs = FakeHeteroDataset(
    num_graphs=num_graphs,
    num_node_types=2,
    num_edge_types=2,
    avg_num_nodes=avg_num_nodes,
    avg_degree=10,
    num_classes=2
)
# collect metadata
graph_features = {}
for node_type, values in vars(graphs[0])['_node_store_dict'].items():
    graph_features[node_type] = len(values['x'][0])
metadata = graphs.data.metadata()
print(metadata)


# train
# 1. train GCN on graph classification
train_dataset, test_dataset = train_test_split(graphs, test_size=0.2)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

model = HGT(hidden_channels=hidden_channels, out_channels=out_channels, num_heads=num_heads, num_layers=num_layers, meta_features=graph_features, metadata=metadata)
#model = HAN(in_channels=-1, out_channels=3, hidden_channels=hidden_channels, heads=num_heads, metadata=metadata, num_classes=num_classes)
# model = GCN(hidden_channels=hidden_channels, out_channels=num_classes)
# model = to_hetero(model, metadata, aggr='sum')
model.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = torch.nn.BCELoss()

current_min_loss = 100000000
patience_counter = 0
for epoch in range(1, 5):
    train(model, train_loader, optimizer, criterion, device)
    train_true, train_pred, train_loss = test(model, train_loader, device)
    test_true, test_pred, test_loss = test(model, test_loader, device)

[rusty1s]

I don't know of any heterogeneous graph classification benchmark dataset, that's why we don't provide an example yet (please ping me if you have any pointers).

I haven't run the code since I'm on phone currently (awesome to see you are already utilizing FakeHeteroDataset), but I think the issue is in the way you do "global pooling". IMO, there exists two options here:

1. You aggregate the global vectors of each node type once again:

   concat = []
   for node_type in self.metadata[0]:
       concat.append(global_mean_pool(out[node_type], batch[node_type]))
   out = torch.stack(concat, dim=0).sum(dim=0)

   This results in an output shape of [num_graphs, hidden_channels].
2. You concatenate them in the feature dimension (not in the batch dimension):

   concat = []
   for node_type in self.metadata[0]:
       concat.append(global_mean_pool(out[node_type], batch[node_type]))
   out = torch.cat(concat, dim=-1)

   This results in an output shape of [num_graphs, num_node_types * hidden_channels].

[carlosSorcero]

Thanks @rusty1s for your quick response!

I think I'm running in circles here. I don't think I want to concatenate in the batch dimension -this would mean that different graphs in the same batch are averaged or influence each other?

Concatenating them in the feature dimension is what I think I want (average / sum all values from features on different node types for the same graph). That brings me to a shape of [num_graphs, num_node_types * hidden_channels]. To go down to 1 value per graph, I am doing:

def forward(self, x_dict, edge_index_dict, batch):
        x = self.han_conv(x_dict, edge_index_dict)

        concat = []
        for node_type in self.metadata[0]:
            concat.append(global_mean_pool(x[node_type], batch[node_type]))
        x = torch.cat(concat, dim=-1)
        
        # Apply a final classifier
        x = F.dropout(x, p=dropout, training=self.training)
        x = self.lin1(x).relu()
        x = torch.sigmoid(self.lin2(x))

        x = torch.reshape(x, (-1,))

        return x

Progress!

After that, I get the following error:

Traceback (most recent call last):
  File "graph_training.py", line 61, in <module>
    train(model, train_loader, optimizer, criterion, device)
  File "/home/paperspace/sorcero/gnn_disambiguation/hetero_training_example/training_utils.py", line 119, in train
    loss = criterion(out, data.y)  # Compute the loss.
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 603, in forward
    return F.binary_cross_entropy(input, target, weight=self.weight, reduction=self.reduction)
  File "/home/paperspace/src/miniconda/envs/pyg/lib/python3.8/site-packages/torch/nn/functional.py", line 2915, in binary_cross_entropy
    return torch._C._nn.binary_cross_entropy(input, target, weight, reduction_enum)
RuntimeError: Found dtype Long but expected Float

Which seems to indicate the data.y is of type Long. Since the type is set by FakeHeteroDataset(), I could workaround by using data.y.float() instead. This seems to let me train!

I'll test with real data and validate a bit further I am aggregating as I expect.

Thank you again

PS: To make it work, I had to also include the number of node types when generating my linear layers:

class HAN(torch.nn.Module):
    def __init__(self, in_channels,
                 out_channels, hidden_channels=128, heads=8, metadata=None, num_classes=1):
        super().__init__()
        self.metadata = metadata
        self.hidden_channels = hidden_channels
        self.han_conv = HANConv(in_channels, self.hidden_channels, heads=heads,
                                dropout=dropout, metadata=metadata)
        self.lin1 = Linear(self.hidden_channels * len(metadata[0]), out_channels)
        self.lin2 = Linear(out_channels, num_classes)

[rusty1s]

Super :)

[ysgncss]

Could you please provide code examples? I have been doing heterogeneous graph classification recently, but there are no relevant documents.

[rusty1s]

We would love to provide a code example, but I am not aware of any publicly available heterogeneous graph classification dataset. Please ping me if you know one. Happy to help you with any questions in order to implement #3758 (comment) though!