init project

Author: psp3dcg

README.md

@@ -1,2 +1,29 @@
 # GSAPool
 Pytorch Implementation of GSAPool (WWW 2020)
+====
+
+PyTorch implementation of [Structure-Feature based Graph Self-adaptive Pooling](https://arxiv.org/pdf/2002.00848)
+
+
+
+
+## Requirements
+  * pytorch
+  * torch_geometric
+
+## Usage
+
+```python control_shell.py```
+
+
+## Cite
+```
+@InProceedings{GSAPool2020,
+  title = 	 {Structure-Feature based Graph Self-adaptive Pooling},
+  author = 	 {Liang Zhang and Xudong Wang and Hongsheng Li and Guangming Zhu and Peiyi Shen and Ping Li and Xiaoyuan Lu and Syed Afaq Ali Shah and  Mohammed Bennamoun},
+  booktitle = {Proceedings of the Web Conference 2020},
+  year = 	 {2020},
+  month = 	 {20-25 April}
+}
+```
+

control_shell.py

@@ -0,0 +1,22 @@
+#coding = "utf-8"
+import os
+import util
+
+parser = util.parser
+args = parser.parse_args()
+os.chdir(args.save_path)
+for i in range(args.training_times):
+    print('------------------------------')
+    print("GSAPool Training Control Shell")
+    print('------------------------------')
+    print('Training Dataset:   ', args.dataset)
+    print('Pooling Layer Type: ',args.pooling_layer_type)
+    print('Feature Fusion Type:',args.feature_fusion_type)
+    print('------------------------------')
+    os.system("python main.py")
+with open(os.path.join(args.save_path, 'result.txt'), 'a') as f:
+    f.write('\r\n')
+    
+
+
+

latest.pth

@@ -0,0 +1,79 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch.nn import Parameter
+from torch_geometric.nn.pool.topk_pool import topk,filter_adj
+from torch_geometric.nn import GCNConv, SAGEConv, GATConv, ChebConv, GraphConv
+
+class GSAPool(torch.nn.Module):
+
+    def __init__(self, in_channels, pooling_ratio=0.5, alpha=0.6, pooling_conv="GCNConv", fusion_conv="false",
+	                    min_score=None, multiplier=1, non_linearity=torch.tanh):
+        super(GSAPool,self).__init__()
+        self.in_channels = in_channels
+		
+        self.ratio = pooling_ratio
+        self.alpha = alpha
+		
+        self.sbtl_layer = self.conv_selection(pooling_conv, in_channels)
+        self.fbtl_layer = nn.Linear(in_channels, 1)
+        self.fusion = self.conv_selection(fusion_conv, in_channels, conv_type=1)
+		
+        self.min_score = min_score
+        self.multiplier = multiplier
+        self.fusion_flag = 0
+        if(fusion_conv!="false"):
+            self.fusion_flag = 1
+        self.non_linearity = non_linearity
+
+    def conv_selection(self, conv, in_channels, conv_type=0):
+        if(conv_type == 0):
+            out_channels = 1
+        elif(conv_type == 1):
+            out_channels = in_channels
+        if(conv == "GCNConv"):
+            return GCNConv(in_channels,out_channels)
+        elif(conv == "ChebConv"):
+            return ChebConv(in_channels,out_channels,1)
+        elif(conv == "SAGEConv"):
+            return SAGEConv(in_channels,out_channels)
+        elif(conv == "GATConv"):
+            return GATConv(in_channels,out_channels, heads=1, concat=True)
+        elif(conv == "GraphConv"):
+            return GraphConv(in_channels,out_channels)
+        else:
+            raise ValueError
+
+    def forward(self, x, edge_index, edge_attr=None, batch=None):
+        if batch is None:
+            batch = edge_index.new_zeros(x.size(0))
+        x = x.unsqueeze(-1) if x.dim() == 1 else x
+
+        #SBTL
+        score_s = self.sbtl_layer(x,edge_index).squeeze()
+        #FBTL
+        score_f = self.fbtl_layer(x).squeeze()
+        #hyperparametr alpha
+        score = score_s*self.alpha + score_f*(1-self.alpha)
+
+        score = score.unsqueeze(-1) if score.dim()==0 else score
+		
+        if self.min_score is None:
+            score = self.non_linearity(score)
+        else:
+            score = softmax(score, batch)
+        perm = topk(score, self.ratio, batch)
+		
+		#fusion
+        if(self.fusion_flag == 1):
+            x = self.fusion(x, edge_index)
+    
+        x = x[perm] * score[perm].view(-1, 1)
+        x = self.multiplier * x if self.multiplier != 1 else x
+        
+        batch = batch[perm]
+        edge_index, edge_attr = filter_adj(
+            edge_index, edge_attr, perm, num_nodes=score.size(0))
+
+        return x, edge_index, edge_attr, batch, perm

layers.py

@@ -0,0 +1,101 @@
+import os
+import torch
+import torch.nn.functional as F
+from torch.utils.data import random_split
+from torch_geometric.datasets import TUDataset
+from torch_geometric.data import DataLoader
+
+
+import util
+from networks import Net
+
+
+
+#parameter initialization
+parser = util.parser
+args = parser.parse_args()
+torch.manual_seed(args.seed)
+
+#device selection
+if torch.cuda.is_available():
+    torch.cuda.manual_seed(args.seed)
+    args.device = 'cuda:0'
+else:
+    args.device = 'cpu'
+
+#dataset split
+def data_builder(args):
+    dataset = TUDataset(os.path.join('data',args.dataset),name=args.dataset)
+    args.num_classes = dataset.num_classes
+    args.num_features = dataset.num_features
+
+    num_training = int(len(dataset)*0.8)
+    num_val = int(len(dataset)*0.1)
+    num_test = len(dataset) - (num_training+num_val)
+    training_set,validation_set,test_set = random_split(dataset,[num_training,num_val,num_test])
+
+    train_loader = DataLoader(training_set, batch_size=args.batch_size, shuffle=True)
+    val_loader = DataLoader(validation_set,batch_size=args.batch_size,shuffle=False)
+    test_loader = DataLoader(test_set,batch_size=1,shuffle=False)
+
+    return train_loader, val_loader, test_loader
+   
+#test function
+def test(model,loader):
+    model.eval()
+    correct = 0.
+    loss = 0.
+    for data in loader:
+        data = data.to(args.device)
+        out = model(data)
+        pred = out.max(dim=1)[1]
+        correct += pred.eq(data.y).sum().item()
+        loss += F.nll_loss(out,data.y,reduction='sum').item()
+    return correct / len(loader.dataset),loss / len(loader.dataset)
+	
+#save result in txt
+def save_result(test_acc, save_path):
+    with open(os.path.join(save_path, 'result.txt'), 'a') as f:
+        test_acc *= 100
+        f.write(args.dataset+";")
+        f.write("pooling_layer_type:"+args.pooling_layer_type+";")
+        f.write("feature_fusion_type:"+args.feature_fusion_type+";")
+        f.write(str(test_acc))
+        f.write('\r\n')
+
+#training configuration
+train_loader, val_loader, test_loader = data_builder(args)
+model = Net(args).to(args.device)
+optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+
+#training steps
+patience = 0
+for epoch in range(args.epochs):
+    model.train()
+    for i, data in enumerate(train_loader):
+        data = data.to(args.device)
+        out = model(data)
+        loss = F.nll_loss(out, data.y)
+        print("Training loss:{}".format(loss.item()))
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+    val_acc,val_loss = test(model,val_loader)
+    print("Validation loss:{}\taccuracy:{}".format(val_loss,val_acc))
+    print("Epoch{}".format(epoch))
+    if val_loss < args.min_loss:
+        torch.save(model.state_dict(),'latest.pth')
+        print("Model saved at epoch{}".format(epoch))
+        min_loss = val_loss
+        patience = 0
+    else:
+        patience += 1
+    if patience > args.patience:
+        break 
+
+#test step
+model = Net(args).to(args.device)
+model.load_state_dict(torch.load('latest.pth'))
+test_acc,test_loss = test(model,test_loader)
+print("Test accuarcy:{}".format(test_acc))
+save_result(test_acc, args.save_path)

main.py

@@ -0,0 +1,66 @@
+import torch
+import numpy as np
+import torch.nn.functional as F
+from torch_geometric.nn import GCNConv
+from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp
+
+from layers import GSAPool
+
+class Net(torch.nn.Module):
+    def __init__(self,args):
+        super(Net, self).__init__()
+		
+        self.args = args
+        self.nhid = args.nhid
+
+        self.num_features = args.num_features
+        self.num_classes = args.num_classes
+        
+        self.alpha = args.alpha
+        self.pooling_ratio = args.pooling_ratio
+        self.dropout_ratio = args.dropout_ratio
+		
+        self.pooling_layer_type = args.pooling_layer_type
+        self.feature_fusion_type = args.feature_fusion_type
+		
+        self.conv1 = GCNConv(self.num_features, self.nhid)
+        self.pool1 = GSAPool(self.nhid, pooling_ratio=self.pooling_ratio, alpha = self.alpha, 
+		                     pooling_conv=self.pooling_layer_type, fusion_conv=self.feature_fusion_type)
+        self.conv2 = GCNConv(self.nhid, self.nhid)
+        self.pool2 = GSAPool(self.nhid, pooling_ratio=self.pooling_ratio, alpha = self.alpha, 
+		                     pooling_conv=self.pooling_layer_type, fusion_conv=self.feature_fusion_type)
+        self.conv3 = GCNConv(self.nhid, self.nhid)
+        self.pool3 = GSAPool(self.nhid, pooling_ratio=self.pooling_ratio, alpha = self.alpha, 
+		                     pooling_conv=self.pooling_layer_type, fusion_conv=self.feature_fusion_type)
+
+        self.lin1 = torch.nn.Linear(self.nhid*2, self.nhid)
+        self.lin2 = torch.nn.Linear(self.nhid, self.nhid//2)
+        self.lin3 = torch.nn.Linear(self.nhid//2, self. num_classes)
+
+  
+    def forward(self, data):
+        x, edge_index, batch = data.x, data.edge_index, data.batch
+
+		
+        x = F.relu(self.conv1(x, edge_index))
+        x, edge_index, _, batch, _ = self.pool1(x, edge_index, None, batch)
+        x1 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
+
+        x = F.relu(self.conv2(x, edge_index))
+        x, edge_index, _, batch, _ = self.pool2(x, edge_index, None, batch)
+        x2 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
+
+        x = F.relu(self.conv3(x, edge_index))
+        x, edge_index, _, batch, _ = self.pool3(x, edge_index, None, batch)
+        x3 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
+
+        x = x1 + x2 + x3
+		
+        x = F.relu(self.lin1(x))
+        x = F.dropout(x, p=self.dropout_ratio, training=self.training)
+        x = F.relu(self.lin2(x))
+        x = F.log_softmax(self.lin3(x), dim=-1)
+
+        return x
+
+    

networks.py

@@ -0,0 +1,40 @@
+import argparse
+#Parameter Configuration
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument('--seed', type=int, default=777,
+                    help='seed')
+parser.add_argument('--batch_size', type=int, default=128,
+                    help='batch size')
+parser.add_argument('--lr', type=float, default=0.0005,
+                    help='learning rate')
+parser.add_argument('--weight_decay', type=float, default=0.0001,
+                    help='weight decay')
+parser.add_argument('--min_loss', type=float, default=1e10,
+                    help='min loss value')
+parser.add_argument('--nhid', type=int, default=128,
+                    help='hidden size')
+parser.add_argument('--pooling_ratio', type=float, default=0.5,
+                    help='pooling ratio')
+parser.add_argument('--alpha', type=float, default=0.6,
+                    help='combination_ratio')
+parser.add_argument('--dropout_ratio', type=float, default=0.5,
+                    help='dropout ratio') 
+parser.add_argument('--dataset', type=str, default='DD',
+                    help='DD/NCI1/NCI109/Mutagenicity')
+parser.add_argument('--epochs', type=int, default=100000,#default = 100000
+                    help='maximum number of epochs')
+parser.add_argument('--patience', type=int, default=50,
+                    help='patience for earlystopping')
+parser.add_argument('--pooling_layer_type', type=str, default='GCNConv',
+                    help='GCNConv')
+parser.add_argument('--feature_fusion_type', type=str, default='GATConv',
+                    help='GCNConv/SAGEConv/ChebConv/GATConv/GraphConv')
+parser.add_argument('--save_path', type=str, default='/home/baoke/workspace_wxd/GSAPool',
+                    help='path to save result')
+parser.add_argument('--training_times', type=int, default=20,
+                    help='')
+
+
+