How to implement ReadOut layer based on torch_geometric ?

[SongZan222]

My job is to complete a graph classification task based on GCN and I want to do it on torch_geometric .

For example, a batch of a training contains G graphs, a total of n nodes. After the node feature matrix is convolved , the node feature matrix of row n needs to be converted into graph feature of row G. The feature of each graph needs to integrate the attributes belonging to the graph node in a specific way, such as taking the maximum value of each component as the value of the graph feature on this component, or directly taking the average value .

The following is a concrete implementation of my model

import numpy as np
import torch
import torch_geometric
from torch_geometric.datasets import TUDataset
from  torch_geometric.loader import DataLoader
from torch_geometric.graphgym import register_pooling
from  torch_geometric.nn import GCNConv
import torch_scatter
import scipy

# Prepare  ENZYMES  data set
enzymes_dataset = TUDataset(
    root=r'C:\Users\Administrator\Desktop\data',
    name='ENZYMES'
)

enzymes_loader = DataLoader(
    batch_size=16,
    shuffle=False
)

# Convolution layer definition
class Net(torch.nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.conv1 = GCNConv(enzymes_dataset.num_node_features,16)
        self.conv2 = GCNConv(16, 32)
        self.fc1 = torch.nn.Linear(32,16)
        self.fc2 = torch.nn.Linear(16,enzymes_dataset.num_classes)

    def forward(self,data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x,edge_index)
        x = torch.nn.functional.relu(x)
        x = self.conv2(x,edge_index)
        x = torch.nn.functional.relu(x)

       # x = Read_Out (x)   This's the function I'm asking for help 

        x = self.fc1(x)
        x = self.fc2(x)
        return torch.nn.functional.log_softmax(x,dim=-1)

[rusty1s]

You can leverage any torch_geoemtric.Aggregation module for the final readout. Here is an example that uses global_add_pool.

[SongZan222]

Dear rusty1s,
Thanks to your help, my model can run smoothly now, but the accuracy is slightly lower than that in the paper. I guess it is because of too few training times.
Here is my model for others to look at.

import numpy as np
import torch
import torch_geometric
from torch_geometric.datasets import TUDataset
from  torch_geometric.loader import DataLoader
from torch_geometric.graphgym import register_pooling
from  torch_geometric.nn import GCNConv
import torch_scatter
import scipy
from torch_geometric.nn import aggr


enzymes_dataset = TUDataset(
    root=r'C:\Users\Administrator\Desktop\data',
    name='ENZYMES'
)

#  spilt dataset 切分数据集
dataset_train = enzymes_dataset[ :500]
dataset_test = enzymes_dataset[500: ]

# Datalorder
def get_dataloader(name):
    dataset = TUDataset(
        root=r'C:\Users\Administrator\Desktop\data',
        name=name
    )
    loader = DataLoader(
        dataset,
        shuffle=True,
        batch_size=64
    )
    return loader


sum_aggr =aggr.SumAggregation()

class Net(torch.nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.conv1 = GCNConv(enzymes_dataset.num_node_features,16)
        self.conv2 = GCNConv(16, 32)
        self.fc1 = torch.nn.Linear(32,16)
        self.fc2 = torch.nn.Linear(16,enzymes_dataset.num_classes)

    def forward(self,data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x,edge_index)
        x = torch.nn.functional.relu(x)
        x = self.conv2(x,edge_index)
        x = torch.nn.functional.relu(x)
        x = sum_aggr(x, ptr=data.ptr)   # Read Out layer
        x = self.fc1(x)
        x = self.fc2(x)

        return torch.nn.functional.log_softmax(x,dim=-1)

model = Net()
optimizer = torch.optim.Adam(model.parameters(),lr=0.01)

# train function 训练函数
def train(epoch):
    data_loder = get_dataloader('ENZYMES')
    counter = 0
    for idx in data_loder:
        output = model(idx).squeeze()
        optimizer.zero_grad()
        loss = torch.nn.functional.nll_loss(output,idx.y)
        loss.backward()
        optimizer.step()
        print(epoch,counter,loss.item())
        counter +=1


for epoch in range(200):
    train(epoch)

def test():
    model.eval()
    dataset = get_dataloader('ENZYMES')
    for batch in dataset:
        out = model(batch).max(1)[1]
        test_acc = torch.eq(
            out, batch.y
        )
        print(out)
        print(test_acc.float().mean())

test()

[Wangyixinxin]

Hi, it seems you didn't use your data splits... you train and test on the same data lol

[SongZan222]

The main goal was to write the body of the model without focusing on the data set, which I later added. You can refer to my question under PyG community. In the answer, the TUDataset class is divided into training set and test set
How to correctly divide the TUdataset class into training sets and test sets ?