├── .gitignore ├── datasets.py ├── README.md ├── tuning.py ├── models ├── sgc.py ├── gfnn.py ├── gcn.py ├── gat.py └── masked_gcn.py ├── LICENSE ├── main.py └── train.py /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | __pycache__/ -------------------------------------------------------------------------------- /datasets.py: -------------------------------------------------------------------------------- 1 | from torch_geometric.datasets import Planetoid 2 | import torch_geometric.transforms as T 3 | 4 | 5 | def get_planetoid_dataset(data_name, normalize_features): 6 | dataset = Planetoid(root='/tmp/' + data_name, name=data_name) 7 | if normalize_features: 8 | dataset.transform = T.NormalizeFeatures() 9 | 10 | return dataset 11 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # GNN_models_pytorch_geometric 2 | 3 | - GCN 4 | - [Semi-Supervised Classification with Graph Convolutional Networks](https://arxiv.org/abs/1609.02907) 5 | - Thomas N. Kipf, Max Welling 6 | - ICLR 2017 7 | - GAT 8 | - [Graph Attention Networks](https://arxiv.org/abs/1710.10903) 9 | - Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, Yoshua Bengio 10 | - ICLR 2018 11 | - SGC 12 | - [Simplifying Graph Convolutional Networks](https://arxiv.org/abs/1902.07153) 13 | - Felix Wu, Tianyi Zhang, Amauri Holanda de Souza Jr., Christopher Fifty, Tao Yu, Kilian Q. Weinberger 14 | - ICML 2019 15 | -------------------------------------------------------------------------------- /tuning.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | from hyperopt import fmin, tpe, hp, STATUS_OK 3 | from math import log 4 | from torch.optim import Adam 5 | 6 | from train import run 7 | 8 | 9 | def objective(dataset, model, lr, space): 10 | optimizer = Adam(model.parameters(), lr=lr, weight_decay=space['weight_decay']) 11 | evals = run(dataset, model, optimizer, early_stopping=False) 12 | return { 13 | 'loss': -evals['val_acc'], 14 | 'status': STATUS_OK 15 | } 16 | 17 | 18 | def search_best_hp(dataset, model, lr): 19 | f = partial(objective, dataset, model, lr) 20 | space = {'weight_decay': hp.loguniform('weight_decay', log(1e-9), log(1e-3))} 21 | best = fmin(fn=f, space=space, algo=tpe.suggest, max_evals=60) 22 | print(best['weight_decay']) 23 | return best['weight_decay'] 24 | -------------------------------------------------------------------------------- /models/sgc.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | from torch.optim import Adam 4 | from torch_geometric.nn import SGConv 5 | 6 | from datasets import get_planetoid_dataset 7 | 8 | 9 | class SGC(nn.Module): 10 | def __init__(self, dataset, K): 11 | super(SGC, self).__init__() 12 | self.gc1 = SGConv(dataset.num_features, dataset.num_classes, K=K, cached=True) 13 | 14 | def reset_parameters(self): 15 | self.gc1.reset_parameters() 16 | 17 | def forward(self, data): 18 | x, edge_index = data.x, data.edge_index 19 | x = self.gc1(x, edge_index) 20 | return F.log_softmax(x, dim=1) 21 | 22 | 23 | def create_sgc_model(data_name, lr=0.2, weight_decay=3e-5): 24 | dataset = get_planetoid_dataset(data_name, True) 25 | model = SGC(dataset, 2) 26 | optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay) 27 | 28 | return dataset, model, optimizer 29 | -------------------------------------------------------------------------------- /models/gfnn.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | from torch.optim import Adam 4 | from torch_geometric.nn import SGConv 5 | 6 | from datasets import get_planetoid_dataset 7 | 8 | 9 | class GFNN(nn.Module): 10 | def __init__(self, dataset, nhid, K): 11 | super(GFNN, self).__init__() 12 | self.gc1 = SGConv(dataset.num_features, nhid, K=K, cached=True) 13 | self.fc1 = nn.Linear(nhid, dataset.num_classes) 14 | 15 | def reset_parameters(self): 16 | self.gc1.reset_parameters() 17 | self.fc1.reset_parameters() 18 | 19 | def forward(self, data): 20 | x, edge_index = data.x, data.edge_index 21 | x = self.gc1(x, edge_index) 22 | x = F.relu(x) 23 | x = self.fc1(x) 24 | return F.log_softmax(x, dim=1) 25 | 26 | 27 | def create_gfnn_model(data_name, nhid=32, lr=0.2, weight_decay=5e-4): 28 | dataset = get_planetoid_dataset(data_name, True) 29 | model = GFNN(dataset, nhid, 2) 30 | optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay) 31 | return dataset, model, optimizer 32 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Hibiki Taguchi 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /models/gcn.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | from torch.optim import Adam 4 | from torch_geometric.nn import GCNConv 5 | 6 | from datasets import get_planetoid_dataset 7 | 8 | 9 | class GCN(nn.Module): 10 | def __init__(self, dataset, nhid, dropout): 11 | super(GCN, self).__init__() 12 | self.gc1 = GCNConv(dataset.num_features, nhid) 13 | self.gc2 = GCNConv(nhid, dataset.num_classes) 14 | self.dropout = dropout 15 | 16 | def reset_parameters(self): 17 | self.gc1.reset_parameters() 18 | self.gc2.reset_parameters() 19 | 20 | def forward(self, data): 21 | x, edge_index = data.x, data.edge_index 22 | x = F.dropout(x, p=self.dropout, training=self.training) 23 | x = self.gc1(x, edge_index) 24 | x = F.relu(x) 25 | x = F.dropout(x, p=self.dropout, training=self.training) 26 | x = self.gc2(x, edge_index) 27 | return F.log_softmax(x, dim=1) 28 | 29 | 30 | def create_gcn_model(data_name, nhid=16, dropout=0.5, 31 | lr=0.01, weight_decay=5e-4): 32 | dataset = get_planetoid_dataset(data_name, True) 33 | model = GCN(dataset, nhid, dropout) 34 | optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay) 35 | 36 | return dataset, model, optimizer 37 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | from models.gcn import create_gcn_model 2 | from models.gat import create_gat_model 3 | from models.gfnn import create_gfnn_model 4 | from models.masked_gcn import create_masked_gcn_model 5 | from models.sgc import create_sgc_model 6 | from train import run 7 | from tuning import search_best_hp 8 | 9 | if __name__ == '__main__': 10 | # GCN 11 | dataset, model, optimizer = create_gcn_model('Cora') 12 | run(dataset, model, optimizer, verbose=False) 13 | 14 | # GAT 15 | # dataset, model, optimizer = create_gat_model('Cora') 16 | # run(dataset, model, optimizer, patience=100) 17 | 18 | # SGC 19 | # Hyper parameter search 20 | # dataset, model, _ = create_sgc_model('Cora') 21 | # weight_decay = search_best_hp(dataset, model, lr=0.2) 22 | # dataset, model, optimizer = create_sgc_model('Cora', weight_decay=weight_decay) 23 | # run(dataset, model, optimizer, epochs=100, early_stopping=False, verbose=False) 24 | 25 | # gfNN 26 | # Hyper parameter search 27 | # dataset, model, _ = create_gfnn_model('Cora') 28 | # weight_decay = search_best_hp(dataset, model, lr=0.2) 29 | # dataset, model, optimizer = create_gfnn_model('Cora', weight_decay=weight_decay) 30 | # run(dataset, model, optimizer, epochs=50, early_stopping=False) 31 | 32 | # Masked GCN 33 | # dataset, model, optimizer = create_masked_gcn_model('Cora') 34 | # run(dataset, model, optimizer, verbose=True) 35 | -------------------------------------------------------------------------------- /models/gat.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | from torch.optim import Adam 4 | from torch_geometric.nn import GATConv 5 | 6 | from datasets import get_planetoid_dataset 7 | 8 | 9 | class GAT(nn.Module): 10 | def __init__(self, dataset, nhid, first_heads, output_heads, dropout): 11 | super(GAT, self).__init__() 12 | self.gc1 = GATConv(dataset.num_features, nhid, 13 | heads=first_heads, dropout=dropout) 14 | self.gc2 = GATConv(nhid*first_heads, dataset.num_classes, 15 | heads=output_heads, dropout=dropout) 16 | self.dropout = dropout 17 | 18 | def reset_parameters(self): 19 | self.gc1.reset_parameters() 20 | self.gc2.reset_parameters() 21 | 22 | def forward(self, data): 23 | x, edge_index = data.x, data.edge_index 24 | x = F.dropout(x, p=self.dropout, training=self.training) 25 | x = self.gc1(x, edge_index) 26 | x = F.elu(x) 27 | x = F.dropout(x, p=self.dropout, training=self.training) 28 | x = self.gc2(x, edge_index) 29 | return F.log_softmax(x, dim=1) 30 | 31 | 32 | def create_gat_model(data_name, nhid=8, first_heads=8, output_heads=1, 33 | dropout=0.6, lr=0.005, weight_decay=5e-4): 34 | dataset = get_planetoid_dataset(data_name, True) 35 | model = GAT(dataset, nhid, first_heads, output_heads, dropout) 36 | optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay) 37 | 38 | return dataset, model, optimizer 39 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | from numpy import mean, std 4 | from tqdm import tqdm 5 | 6 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 7 | 8 | 9 | def train(model, optimizer, data): 10 | model.train() 11 | optimizer.zero_grad() 12 | output = model(data) 13 | loss = F.nll_loss(output[data.train_mask == 1], data.y[data.train_mask == 1]) 14 | loss.backward() 15 | optimizer.step() 16 | 17 | 18 | def evaluate(model, data): 19 | model.eval() 20 | 21 | with torch.no_grad(): 22 | output = model(data) 23 | 24 | outputs = {} 25 | for key in ['train', 'val', 'test']: 26 | mask = data['{}_mask'.format(key)] 27 | loss = F.nll_loss(output[mask == 1], data.y[mask == 1]).item() 28 | pred = output[mask == 1].max(dim=1)[1] 29 | acc = pred.eq(data.y[mask == 1]).sum().item() / mask.sum().item() 30 | 31 | outputs['{}_loss'.format(key)] = loss 32 | outputs['{}_acc'.format(key)] = acc 33 | 34 | return outputs 35 | 36 | 37 | def run(dataset, model, optimizer, epochs=200, iter=100, early_stopping=True, patience=10, verbose=False): 38 | data = dataset[0] 39 | 40 | # for GPU 41 | data = data.to(device) 42 | 43 | if torch.cuda.is_available(): 44 | torch.cuda.synchronize() 45 | 46 | val_acc_list = [] 47 | test_acc_list = [] 48 | 49 | for _ in tqdm(range(iter)): 50 | # for early stopping 51 | model.to(device).reset_parameters() 52 | best_val_loss = float('inf') 53 | counter = 0 54 | for epoch in range(1, epochs+1): 55 | train(model, optimizer, data) 56 | evals = evaluate(model, data) 57 | 58 | if verbose: 59 | print('epoch:', epoch, 'train loss:', evals['train_loss'], 60 | 'val loss:', evals['val_loss']) 61 | 62 | if early_stopping: 63 | if evals['val_loss'] < best_val_loss: 64 | best_val_loss = evals['val_loss'] 65 | counter = 0 66 | else: 67 | counter += 1 68 | if counter >= patience: 69 | # print("Stop training, epoch:", epoch) 70 | break 71 | if verbose: 72 | for met, val in evals.items(): 73 | print(met, val) 74 | 75 | val_acc_list.append(evals['val_acc']) 76 | test_acc_list.append(evals['test_acc']) 77 | 78 | print(mean(test_acc_list)) 79 | print(std(test_acc_list)) 80 | return { 81 | 'val_acc': mean(val_acc_list), 82 | 'test_acc': mean(test_acc_list), 83 | 'test_acc_std': std(test_acc_list) 84 | } 85 | -------------------------------------------------------------------------------- /models/masked_gcn.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from torch.nn import Parameter 6 | from torch.optim import Adam 7 | from torch_geometric.nn import GATConv, GCNConv 8 | from torch_geometric.utils import add_self_loops, degree 9 | 10 | from datasets import get_planetoid_dataset 11 | 12 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 13 | 14 | 15 | def mask_features(x, edge_index, edge_weight, sigma): 16 | source, target = edge_index 17 | h_s, h_t = x[source], x[target] 18 | h = (h_t - h_s) / sigma 19 | h = edge_weight.view(-1, 1) * h * h 20 | mask = torch.zeros(x.size(), device=device) 21 | mask.index_add_(0, source, h) 22 | deg = degree(edge_index[0]) 23 | mask = torch.exp(- mask / deg.view(-1, 1)) 24 | x = x * mask 25 | return x 26 | 27 | 28 | class MaskedGCNConv(GCNConv): 29 | def __init__(self, in_channels, out_channels): 30 | super(MaskedGCNConv, self).__init__(in_channels, out_channels) 31 | self.sigma = Parameter(torch.Tensor(1, in_channels)) 32 | nn.init.xavier_uniform_(self.sigma.data, gain=1.414) 33 | 34 | def reset_parameters(self): 35 | super().reset_parameters() 36 | 37 | def forward(self, x, edge_index, edge_weight=None): 38 | if not self.cached or self.cached_result is None: 39 | self.cached_num_edges = edge_index.size(1) 40 | edge_index, norm = self.norm(edge_index, x.size(0), edge_weight, 41 | self.improved, x.dtype) 42 | self.cached_result = edge_index, norm 43 | edge_index, norm = self.cached_result 44 | x = mask_features(x, edge_index, norm, self.sigma) 45 | ret = super().forward(x, edge_index, edge_weight) 46 | return ret 47 | 48 | 49 | class MaskedGCN(nn.Module): 50 | def __init__(self, dataset, nhid, dropout): 51 | super(MaskedGCN, self).__init__() 52 | self.gc1 = MaskedGCNConv(dataset.num_features, nhid) 53 | self.gc2 = MaskedGCNConv(nhid, dataset.num_classes) 54 | self.dropout = dropout 55 | 56 | def reset_parameters(self): 57 | self.gc1.reset_parameters() 58 | self.gc2.reset_parameters() 59 | 60 | def forward(self, data): 61 | x, edge_index = data.x, data.edge_index 62 | x = F.dropout(x, p=self.dropout, training=self.training) 63 | x = self.gc1(x, edge_index) 64 | x = F.relu(x) 65 | x = F.dropout(x, p=self.dropout, training=self.training) 66 | x = self.gc2(x, edge_index) 67 | return F.log_softmax(x, dim=1) 68 | 69 | 70 | def create_masked_gcn_model(data_name, nhid=16, dropout=0.5, 71 | lr=0.02, weight_decay=5e-4): 72 | dataset = get_planetoid_dataset(data_name, True) 73 | model = MaskedGCN(dataset, nhid, dropout) 74 | optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay) 75 | 76 | return dataset, model, optimizer --------------------------------------------------------------------------------