How to Use PANPooling on TUDatasets

[SongZan222]

For now I want to implement a simple graph neural network aiming at graph classification, using either GCN or GIN for the convolutional layer, I use GIN for the example below. I have successfully used SAGPool in the pooling layer before, but I had trouble extending it to PANPool.
My question is : How to use the PANPool module correctly?
If there is an implementation using ASAPpool mode can also work.
Here are the key parts of my code:

from torch_geometric.nn import GIN
from torch_geometric.nn.pool import SAGPooling
from torch_geometric.nn.pool import PANPooling
from torch_geometric.nn.pool import ASAPooling


# datasets
dataset = TUDataset(
    root=r'C:\Users\Administrator\Desktop\data',
    name=DATASET_NAME,
    transform=torch_geometric.transforms.ToDevice(DEVICE)
)

# ReadOut-method
sum_aggr = aggr.SumAggregation().to(DEVICE)

class Net_SAG(torch.nn.Module):
    def __init__(self):
        super(Net_SAG, self).__init__()
        self.conv1 = GIN(
            in_channels = dataset.num_node_features,
            hidden_channels = HIDDEN_DIMENSION,
            out_channels = HIDDEN_DIMENSION,
            num_layers = 3,
            dropout = DROP_OUT
        )
        self.pool = SAGPooling(HIDDEN_DIMENSION, 0.75)
        self.fc1 = torch.nn.Linear(HIDDEN_DIMENSION, 16)
        self.fc2 = torch.nn.Linear(16, 16)
        self.fc3 = torch.nn.Linear(16, dataset.num_classes)

    def forward(self, data):
        x = data.x
        edge_index = data.edge_index
        x = self.conv1(x, edge_index)
        x = torch.nn.functional.relu(x)
        x,edge_index,_,batch,_,_ = self.pool(
            x=x,
            edge_index=edge_index,
            batch=data.batch
        )
        x = sum_aggr(
            x=x,
            index= batch
        )
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return torch.nn.functional.log_softmax(x, dim=-1)

[rusty1s]

You are right that we are currently missing an example for this, really sorry. However, we show-case how to use this module in https://github.com/pyg-team/pytorch_geometric/blob/master/test/nn/pool/test_pan_pool.py, which should hopefully unblock you. Let me know if you have any doubts.

[SongZan222]

Dear PyG community:
Thank you for your reference materials and help.
I have written an operational PANPool pooling model, but due to my poor mathematical foundation, I haven't understood the convolution and pooling formulas in the original paper.

HIDDEN_DIMENSION = 16
avg_aggr = aggr.MeanAggregation()


class Net_PAN(torch.nn.Module):
    def __init__(self):
        super(Net_PAN, self).__init__()

        self.conv1 = PANConv(dataset.num_node_features, HIDDEN_DIMENSION,0)
        self.conv2 = PANConv(HIDDEN_DIMENSION, HIDDEN_DIMENSION,0)
        self.conv3 = PANConv(HIDDEN_DIMENSION, HIDDEN_DIMENSION,0)
        self.pool1 = PANPooling(HIDDEN_DIMENSION,ratio=0.5)
        self.fc1 = torch.nn.Linear(HIDDEN_DIMENSION,HIDDEN_DIMENSION)
        self.fc2 = torch.nn.Linear(HIDDEN_DIMENSION,dataset.num_classes)
    def forward(self, data):
        x = data.x
        edge_index = data.edge_index
        x, M= self.conv1(x, edge_index)
        x, M= self.conv2(x, edge_index)
        x, M = self.conv3(x, edge_index)
        x, edge_index, edge_weight, batch, perm, score = self.pool1(x,M,data.batch)
        y =avg_aggr(
            x,
            index=batch,
        )
        y = self.fc1(y)
        y = self.fc2(y)
        return torch.nn.functional.log_softmax(y, dim=-1)