Unverified Commit f19f05ce authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] Black auto fix. (#4651)


Co-authored-by: default avatarSteve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
parent 977b1ba4
from torch.utils.data import Dataset, DataLoader import collections
from torch.utils.data import DataLoader, Dataset
import dgl import dgl
from dgl.data import PPIDataset from dgl.data import PPIDataset
import collections
#implement the collate_fn for dgl graph data class # implement the collate_fn for dgl graph data class
PPIBatch = collections.namedtuple('PPIBatch', ['graph', 'label']) PPIBatch = collections.namedtuple("PPIBatch", ["graph", "label"])
def batcher(device): def batcher(device):
def batcher_dev(batch): def batcher_dev(batch):
batch_graphs = dgl.batch(batch) batch_graphs = dgl.batch(batch)
return PPIBatch(graph=batch_graphs, return PPIBatch(
label=batch_graphs.ndata['label'].to(device)) graph=batch_graphs, label=batch_graphs.ndata["label"].to(device)
)
return batcher_dev return batcher_dev
#add a fresh "self-loop" edge type to the untyped PPI dataset and prepare train, val, test loaders
def load_PPI(batch_size=1, device='cpu'): # add a fresh "self-loop" edge type to the untyped PPI dataset and prepare train, val, test loaders
train_set = PPIDataset(mode='train') def load_PPI(batch_size=1, device="cpu"):
valid_set = PPIDataset(mode='valid') train_set = PPIDataset(mode="train")
test_set = PPIDataset(mode='test') valid_set = PPIDataset(mode="valid")
#for each graph, add self-loops as a new relation type test_set = PPIDataset(mode="test")
#here we reconstruct the graph since the schema of a heterograph cannot be changed once constructed # for each graph, add self-loops as a new relation type
# here we reconstruct the graph since the schema of a heterograph cannot be changed once constructed
for i in range(len(train_set)): for i in range(len(train_set)):
g = dgl.heterograph({ g = dgl.heterograph(
('_N','_E','_N'): train_set[i].edges(), {
('_N', 'self', '_N'): (train_set[i].nodes(), train_set[i].nodes()) ("_N", "_E", "_N"): train_set[i].edges(),
}) ("_N", "self", "_N"): (
g.ndata['label'] = train_set[i].ndata['label'] train_set[i].nodes(),
g.ndata['feat'] = train_set[i].ndata['feat'] train_set[i].nodes(),
g.ndata['_ID'] = train_set[i].ndata['_ID'] ),
g.edges['_E'].data['_ID'] = train_set[i].edata['_ID'] }
)
g.ndata["label"] = train_set[i].ndata["label"]
g.ndata["feat"] = train_set[i].ndata["feat"]
g.ndata["_ID"] = train_set[i].ndata["_ID"]
g.edges["_E"].data["_ID"] = train_set[i].edata["_ID"]
train_set.graphs[i] = g train_set.graphs[i] = g
for i in range(len(valid_set)): for i in range(len(valid_set)):
g = dgl.heterograph({ g = dgl.heterograph(
('_N','_E','_N'): valid_set[i].edges(), {
('_N', 'self', '_N'): (valid_set[i].nodes(), valid_set[i].nodes()) ("_N", "_E", "_N"): valid_set[i].edges(),
}) ("_N", "self", "_N"): (
g.ndata['label'] = valid_set[i].ndata['label'] valid_set[i].nodes(),
g.ndata['feat'] = valid_set[i].ndata['feat'] valid_set[i].nodes(),
g.ndata['_ID'] = valid_set[i].ndata['_ID'] ),
g.edges['_E'].data['_ID'] = valid_set[i].edata['_ID'] }
valid_set.graphs[i] = g )
g.ndata["label"] = valid_set[i].ndata["label"]
g.ndata["feat"] = valid_set[i].ndata["feat"]
g.ndata["_ID"] = valid_set[i].ndata["_ID"]
g.edges["_E"].data["_ID"] = valid_set[i].edata["_ID"]
valid_set.graphs[i] = g
for i in range(len(test_set)): for i in range(len(test_set)):
g = dgl.heterograph({ g = dgl.heterograph(
('_N','_E','_N'): test_set[i].edges(), {
('_N', 'self', '_N'): (test_set[i].nodes(), test_set[i].nodes()) ("_N", "_E", "_N"): test_set[i].edges(),
}) ("_N", "self", "_N"): (
g.ndata['label'] = test_set[i].ndata['label'] test_set[i].nodes(),
g.ndata['feat'] = test_set[i].ndata['feat'] test_set[i].nodes(),
g.ndata['_ID'] = test_set[i].ndata['_ID'] ),
g.edges['_E'].data['_ID'] = test_set[i].edata['_ID'] }
test_set.graphs[i] = g )
g.ndata["label"] = test_set[i].ndata["label"]
g.ndata["feat"] = test_set[i].ndata["feat"]
g.ndata["_ID"] = test_set[i].ndata["_ID"]
g.edges["_E"].data["_ID"] = test_set[i].edata["_ID"]
test_set.graphs[i] = g
etypes = train_set[0].etypes etypes = train_set[0].etypes
in_size = train_set[0].ndata['feat'].shape[1] in_size = train_set[0].ndata["feat"].shape[1]
out_size = train_set[0].ndata['label'].shape[1] out_size = train_set[0].ndata["label"].shape[1]
#prepare train, valid, and test dataloaders # prepare train, valid, and test dataloaders
train_loader = DataLoader(train_set, batch_size=batch_size, collate_fn=batcher(device), shuffle=True) train_loader = DataLoader(
valid_loader = DataLoader(valid_set, batch_size=batch_size, collate_fn=batcher(device), shuffle=True) train_set,
test_loader = DataLoader(test_set, batch_size=batch_size, collate_fn=batcher(device), shuffle=True) batch_size=batch_size,
collate_fn=batcher(device),
shuffle=True,
)
valid_loader = DataLoader(
valid_set,
batch_size=batch_size,
collate_fn=batcher(device),
shuffle=True,
)
test_loader = DataLoader(
test_set,
batch_size=batch_size,
collate_fn=batcher(device),
shuffle=True,
)
return train_loader, valid_loader, test_loader, etypes, in_size, out_size return train_loader, valid_loader, test_loader, etypes, in_size, out_size
import argparse
import os
import numpy as np
import torch import torch
import torch.nn as nn import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
from data_loader import load_PPI
from utils import evaluate_f1_score
import dgl import dgl
import dgl.function as fn import dgl.function as fn
from utils import evaluate_f1_score
from data_loader import load_PPI
import argparse
import numpy as np
import os
class GNNFiLMLayer(nn.Module): class GNNFiLMLayer(nn.Module):
def __init__(self, in_size, out_size, etypes, dropout=0.1): def __init__(self, in_size, out_size, etypes, dropout=0.1):
...@@ -15,78 +18,112 @@ class GNNFiLMLayer(nn.Module): ...@@ -15,78 +18,112 @@ class GNNFiLMLayer(nn.Module):
self.in_size = in_size self.in_size = in_size
self.out_size = out_size self.out_size = out_size
#weights for different types of edges # weights for different types of edges
self.W = nn.ModuleDict({ self.W = nn.ModuleDict(
name : nn.Linear(in_size, out_size, bias = False) for name in etypes {name: nn.Linear(in_size, out_size, bias=False) for name in etypes}
}) )
#hypernets to learn the affine functions for different types of edges # hypernets to learn the affine functions for different types of edges
self.film = nn.ModuleDict({ self.film = nn.ModuleDict(
name : nn.Linear(in_size, 2*out_size, bias = False) for name in etypes {
}) name: nn.Linear(in_size, 2 * out_size, bias=False)
for name in etypes
}
)
#layernorm before each propogation # layernorm before each propogation
self.layernorm = nn.LayerNorm(out_size) self.layernorm = nn.LayerNorm(out_size)
#dropout layer # dropout layer
self.dropout = nn.Dropout(dropout) self.dropout = nn.Dropout(dropout)
def forward(self, g, feat_dict): def forward(self, g, feat_dict):
#the input graph is a multi-relational graph, so treated as hetero-graph. # the input graph is a multi-relational graph, so treated as hetero-graph.
funcs = {} #message and reduce functions dict funcs = {} # message and reduce functions dict
#for each type of edges, compute messages and reduce them all # for each type of edges, compute messages and reduce them all
for srctype, etype, dsttype in g.canonical_etypes: for srctype, etype, dsttype in g.canonical_etypes:
messages = self.W[etype](feat_dict[srctype]) #apply W_l on src feature messages = self.W[etype](
film_weights = self.film[etype](feat_dict[dsttype]) #use dst feature to compute affine function paras feat_dict[srctype]
gamma = film_weights[:,:self.out_size] #"gamma" for the affine function ) # apply W_l on src feature
beta = film_weights[:,self.out_size:] #"beta" for the affine function film_weights = self.film[etype](
messages = gamma * messages + beta #compute messages feat_dict[dsttype]
messages = F.relu_(messages) ) # use dst feature to compute affine function paras
g.nodes[srctype].data[etype] = messages #store in ndata gamma = film_weights[
funcs[etype] = (fn.copy_u(etype, 'm'), fn.sum('m', 'h')) #define message and reduce functions :, : self.out_size
g.multi_update_all(funcs, 'sum') #update all, reduce by first type-wisely then across different types ] # "gamma" for the affine function
feat_dict={} beta = film_weights[
:, self.out_size :
] # "beta" for the affine function
messages = gamma * messages + beta # compute messages
messages = F.relu_(messages)
g.nodes[srctype].data[etype] = messages # store in ndata
funcs[etype] = (
fn.copy_u(etype, "m"),
fn.sum("m", "h"),
) # define message and reduce functions
g.multi_update_all(
funcs, "sum"
) # update all, reduce by first type-wisely then across different types
feat_dict = {}
for ntype in g.ntypes: for ntype in g.ntypes:
feat_dict[ntype] = self.dropout(self.layernorm(g.nodes[ntype].data['h'])) #apply layernorm and dropout feat_dict[ntype] = self.dropout(
self.layernorm(g.nodes[ntype].data["h"])
) # apply layernorm and dropout
return feat_dict return feat_dict
class GNNFiLM(nn.Module): class GNNFiLM(nn.Module):
def __init__(self, etypes, in_size, hidden_size, out_size, num_layers, dropout=0.1): def __init__(
self, etypes, in_size, hidden_size, out_size, num_layers, dropout=0.1
):
super(GNNFiLM, self).__init__() super(GNNFiLM, self).__init__()
self.film_layers = nn.ModuleList() self.film_layers = nn.ModuleList()
self.film_layers.append( self.film_layers.append(
GNNFiLMLayer(in_size, hidden_size, etypes, dropout) GNNFiLMLayer(in_size, hidden_size, etypes, dropout)
) )
for i in range(num_layers-1): for i in range(num_layers - 1):
self.film_layers.append( self.film_layers.append(
GNNFiLMLayer(hidden_size, hidden_size, etypes, dropout) GNNFiLMLayer(hidden_size, hidden_size, etypes, dropout)
) )
self.predict = nn.Linear(hidden_size, out_size, bias = True) self.predict = nn.Linear(hidden_size, out_size, bias=True)
def forward(self, g, out_key): def forward(self, g, out_key):
h_dict = {ntype : g.nodes[ntype].data['feat'] for ntype in g.ntypes} #prepare input feature dict h_dict = {
ntype: g.nodes[ntype].data["feat"] for ntype in g.ntypes
} # prepare input feature dict
for layer in self.film_layers: for layer in self.film_layers:
h_dict = layer(g, h_dict) h_dict = layer(g, h_dict)
h = self.predict(h_dict[out_key]) #use the final embed to predict, out_size = num_classes h = self.predict(
h_dict[out_key]
) # use the final embed to predict, out_size = num_classes
h = torch.sigmoid(h) h = torch.sigmoid(h)
return h return h
def main(args): def main(args):
# Step 1: Prepare graph data and retrieve train/validation/test dataloader ============================= # # Step 1: Prepare graph data and retrieve train/validation/test dataloader ============================= #
if args.gpu >= 0 and torch.cuda.is_available(): if args.gpu >= 0 and torch.cuda.is_available():
device = 'cuda:{}'.format(args.gpu) device = "cuda:{}".format(args.gpu)
else: else:
device = 'cpu' device = "cpu"
if args.dataset == 'PPI': if args.dataset == "PPI":
train_set, valid_set, test_set, etypes, in_size, out_size = load_PPI(args.batch_size, device) train_set, valid_set, test_set, etypes, in_size, out_size = load_PPI(
args.batch_size, device
)
# Step 2: Create model and training components=========================================================== # # Step 2: Create model and training components=========================================================== #
model = GNNFiLM(etypes, in_size, args.hidden_size, out_size, args.num_layers).to(device) model = GNNFiLM(
etypes, in_size, args.hidden_size, out_size, args.num_layers
).to(device)
criterion = nn.BCELoss() criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd) optimizer = torch.optim.Adam(
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.step_size, gamma=args.gamma) model.parameters(), lr=args.lr, weight_decay=args.wd
)
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer, args.step_size, gamma=args.gamma
)
# Step 4: training epoches ============================================================================== # # Step 4: training epoches ============================================================================== #
lastf1 = 0 lastf1 = 0
...@@ -101,10 +138,12 @@ def main(args): ...@@ -101,10 +138,12 @@ def main(args):
for batch in train_set: for batch in train_set:
g = batch.graph g = batch.graph
g = g.to(device) g = g.to(device)
logits = model.forward(g, '_N') logits = model.forward(g, "_N")
labels = batch.label labels = batch.label
loss = criterion(logits, labels) loss = criterion(logits, labels)
f1 = evaluate_f1_score(logits.detach().cpu().numpy(), labels.detach().cpu().numpy()) f1 = evaluate_f1_score(
logits.detach().cpu().numpy(), labels.detach().cpu().numpy()
)
optimizer.zero_grad() optimizer.zero_grad()
loss.backward() loss.backward()
...@@ -112,7 +151,7 @@ def main(args): ...@@ -112,7 +151,7 @@ def main(args):
train_loss.append(loss.item()) train_loss.append(loss.item())
train_f1.append(f1) train_f1.append(f1)
train_loss = np.mean(train_loss) train_loss = np.mean(train_loss)
train_f1 = np.mean(train_f1) train_f1 = np.mean(train_f1)
scheduler.step() scheduler.step()
...@@ -121,24 +160,32 @@ def main(args): ...@@ -121,24 +160,32 @@ def main(args):
for batch in valid_set: for batch in valid_set:
g = batch.graph g = batch.graph
g = g.to(device) g = g.to(device)
logits = model.forward(g, '_N') logits = model.forward(g, "_N")
labels = batch.label labels = batch.label
loss = criterion(logits, labels) loss = criterion(logits, labels)
f1 = evaluate_f1_score(logits.detach().cpu().numpy(), labels.detach().cpu().numpy()) f1 = evaluate_f1_score(
logits.detach().cpu().numpy(), labels.detach().cpu().numpy()
)
val_loss.append(loss.item()) val_loss.append(loss.item())
val_f1.append(f1) val_f1.append(f1)
val_loss = np.mean(val_loss) val_loss = np.mean(val_loss)
val_f1 = np.mean(val_f1) val_f1 = np.mean(val_f1)
print('Epoch {:d} | Train Loss {:.4f} | Train F1 {:.4f} | Val Loss {:.4f} | Val F1 {:.4f} |'.format(epoch + 1, train_loss, train_f1, val_loss, val_f1)) print(
"Epoch {:d} | Train Loss {:.4f} | Train F1 {:.4f} | Val Loss {:.4f} | Val F1 {:.4f} |".format(
epoch + 1, train_loss, train_f1, val_loss, val_f1
)
)
if val_f1 > best_val_f1: if val_f1 > best_val_f1:
best_val_f1 = val_f1 best_val_f1 = val_f1
torch.save(model.state_dict(), os.path.join(args.save_dir, args.name)) torch.save(
model.state_dict(), os.path.join(args.save_dir, args.name)
)
if val_f1 < lastf1: if val_f1 < lastf1:
cnt += 1 cnt += 1
if cnt == args.early_stopping: if cnt == args.early_stopping:
print('Early stop.') print("Early stop.")
break break
else: else:
cnt = 0 cnt = 0
...@@ -152,41 +199,89 @@ def main(args): ...@@ -152,41 +199,89 @@ def main(args):
for batch in test_set: for batch in test_set:
g = batch.graph g = batch.graph
g = g.to(device) g = g.to(device)
logits = model.forward(g, '_N') logits = model.forward(g, "_N")
labels = batch.label labels = batch.label
loss = criterion(logits, labels) loss = criterion(logits, labels)
f1 = evaluate_f1_score(logits.detach().cpu().numpy(), labels.detach().cpu().numpy()) f1 = evaluate_f1_score(
logits.detach().cpu().numpy(), labels.detach().cpu().numpy()
)
test_loss.append(loss.item()) test_loss.append(loss.item())
test_f1.append(f1) test_f1.append(f1)
test_loss = np.mean(test_loss) test_loss = np.mean(test_loss)
test_f1 = np.mean(test_f1) test_f1 = np.mean(test_f1)
print("Test F1: {:.4f} | Test loss: {:.4f}".format(test_f1, test_loss)) print("Test F1: {:.4f} | Test loss: {:.4f}".format(test_f1, test_loss))
if __name__ == "__main__":
if __name__ == '__main__': parser = argparse.ArgumentParser(description="GNN-FiLM")
parser = argparse.ArgumentParser(description='GNN-FiLM') parser.add_argument(
parser.add_argument("--dataset", type=str, default="PPI", help="DGL dataset for this GNN-FiLM") "--dataset",
parser.add_argument("--gpu", type=int, default=-1, help="GPU Index. Default: -1, using CPU.") type=str,
parser.add_argument("--in_size", type=int, default=50, help="Input dimensionalities") default="PPI",
parser.add_argument("--hidden_size", type=int, default=320, help="Hidden layer dimensionalities") help="DGL dataset for this GNN-FiLM",
parser.add_argument("--out_size", type=int, default=121, help="Output dimensionalities") )
parser.add_argument("--num_layers", type=int, default=4, help="Number of GNN layers") parser.add_argument(
"--gpu", type=int, default=-1, help="GPU Index. Default: -1, using CPU."
)
parser.add_argument(
"--in_size", type=int, default=50, help="Input dimensionalities"
)
parser.add_argument(
"--hidden_size",
type=int,
default=320,
help="Hidden layer dimensionalities",
)
parser.add_argument(
"--out_size", type=int, default=121, help="Output dimensionalities"
)
parser.add_argument(
"--num_layers", type=int, default=4, help="Number of GNN layers"
)
parser.add_argument("--batch_size", type=int, default=5, help="Batch size") parser.add_argument("--batch_size", type=int, default=5, help="Batch size")
parser.add_argument("--max_epoch", type=int, default=1500, help="The max number of epoches. Default: 500") parser.add_argument(
parser.add_argument("--early_stopping", type=int, default=80, help="Early stopping. Default: 50") "--max_epoch",
parser.add_argument("--lr", type=float, default=0.001, help="Learning rate. Default: 3e-1") type=int,
parser.add_argument("--wd", type=float, default=0.0009, help="Weight decay. Default: 3e-1") default=1500,
parser.add_argument('--step-size', type=int, default=40, help='Period of learning rate decay.') help="The max number of epoches. Default: 500",
parser.add_argument('--gamma', type=float, default=0.8, help='Multiplicative factor of learning rate decay.') )
parser.add_argument("--dropout", type=float, default=0.1, help="Dropout rate. Default: 0.9") parser.add_argument(
parser.add_argument('--save_dir', type=str, default='./out', help='Path to save the model.') "--early_stopping",
parser.add_argument("--name", type=str, default='GNN-FiLM', help="Saved model name.") type=int,
default=80,
help="Early stopping. Default: 50",
)
parser.add_argument(
"--lr", type=float, default=0.001, help="Learning rate. Default: 3e-1"
)
parser.add_argument(
"--wd", type=float, default=0.0009, help="Weight decay. Default: 3e-1"
)
parser.add_argument(
"--step-size",
type=int,
default=40,
help="Period of learning rate decay.",
)
parser.add_argument(
"--gamma",
type=float,
default=0.8,
help="Multiplicative factor of learning rate decay.",
)
parser.add_argument(
"--dropout", type=float, default=0.1, help="Dropout rate. Default: 0.9"
)
parser.add_argument(
"--save_dir", type=str, default="./out", help="Path to save the model."
)
parser.add_argument(
"--name", type=str, default="GNN-FiLM", help="Saved model name."
)
args = parser.parse_args() args = parser.parse_args()
print(args) print(args)
if not os.path.exists(args.save_dir): if not os.path.exists(args.save_dir):
os.mkdir(args.save_dir) os.mkdir(args.save_dir)
main(args) main(args)
from sklearn.metrics import f1_score
import numpy as np import numpy as np
from sklearn.metrics import f1_score
#function to compute f1 score # function to compute f1 score
def evaluate_f1_score(pred, label): def evaluate_f1_score(pred, label):
pred = np.round(pred, 0).astype(np.int16) pred = np.round(pred, 0).astype(np.int16)
pred = pred.flatten() pred = pred.flatten()
......
...@@ -3,49 +3,100 @@ Data utils for processing bAbI datasets ...@@ -3,49 +3,100 @@ Data utils for processing bAbI datasets
""" """
import os import os
import string
import torch
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
import dgl import dgl
import torch from dgl.data.utils import (
import string _get_dgl_url,
from dgl.data.utils import download, get_download_dir, _get_dgl_url, extract_archive download,
extract_archive,
get_download_dir,
)
def get_babi_dataloaders(batch_size, train_size=50, task_id=4, q_type=0): def get_babi_dataloaders(batch_size, train_size=50, task_id=4, q_type=0):
_download_babi_data() _download_babi_data()
node_dict = dict(zip(list(string.ascii_uppercase), range(len(string.ascii_uppercase)))) node_dict = dict(
zip(list(string.ascii_uppercase), range(len(string.ascii_uppercase)))
)
if task_id == 4: if task_id == 4:
edge_dict = {'n': 0, 's': 1, 'w': 2, 'e': 3} edge_dict = {"n": 0, "s": 1, "w": 2, "e": 3}
reverse_edge = {} reverse_edge = {}
return _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse_edge, '04') return _ns_dataloader(
train_size,
q_type,
batch_size,
node_dict,
edge_dict,
reverse_edge,
"04",
)
elif task_id == 15: elif task_id == 15:
edge_dict = {'is': 0, 'has_fear': 1} edge_dict = {"is": 0, "has_fear": 1}
reverse_edge = {} reverse_edge = {}
return _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse_edge, '15') return _ns_dataloader(
train_size,
q_type,
batch_size,
node_dict,
edge_dict,
reverse_edge,
"15",
)
elif task_id == 16: elif task_id == 16:
edge_dict = {'is': 0, 'has_color': 1} edge_dict = {"is": 0, "has_color": 1}
reverse_edge = {0: 0} reverse_edge = {0: 0}
return _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse_edge, '16') return _ns_dataloader(
train_size,
q_type,
batch_size,
node_dict,
edge_dict,
reverse_edge,
"16",
)
elif task_id == 18: elif task_id == 18:
edge_dict = {'>': 0, '<': 1} edge_dict = {">": 0, "<": 1}
label_dict = {'false': 0, 'true': 1} label_dict = {"false": 0, "true": 1}
reverse_edge = {0: 1, 1: 0} reverse_edge = {0: 1, 1: 0}
return _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_dict, reverse_edge, '18') return _gc_dataloader(
train_size,
q_type,
batch_size,
node_dict,
edge_dict,
label_dict,
reverse_edge,
"18",
)
elif task_id == 19: elif task_id == 19:
edge_dict = {'n': 0, 's': 1, 'w': 2, 'e': 3, '<end>': 4} edge_dict = {"n": 0, "s": 1, "w": 2, "e": 3, "<end>": 4}
reverse_edge = {0: 1, 1: 0, 2: 3, 3: 2} reverse_edge = {0: 1, 1: 0, 2: 3, 3: 2}
max_seq_length = 2 max_seq_length = 2
return _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, reverse_edge, '19', max_seq_length) return _path_finding_dataloader(
train_size,
batch_size,
def _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse_edge, path): node_dict,
edge_dict,
reverse_edge,
"19",
max_seq_length,
)
def _ns_dataloader(
train_size, q_type, batch_size, node_dict, edge_dict, reverse_edge, path
):
def _collate_fn(batch): def _collate_fn(batch):
graphs = [] graphs = []
labels = [] labels = []
for d in batch: for d in batch:
edges = d['edges'] edges = d["edges"]
node_ids = [] node_ids = []
for s, e, t in edges: for s, e, t in edges:
...@@ -55,12 +106,12 @@ def _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse ...@@ -55,12 +106,12 @@ def _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse
node_ids.append(t) node_ids.append(t)
g = dgl.DGLGraph() g = dgl.DGLGraph()
g.add_nodes(len(node_ids)) g.add_nodes(len(node_ids))
g.ndata['node_id'] = torch.tensor(node_ids, dtype=torch.long) g.ndata["node_id"] = torch.tensor(node_ids, dtype=torch.long)
nid2idx = dict(zip(node_ids, list(range(len(node_ids))))) nid2idx = dict(zip(node_ids, list(range(len(node_ids)))))
# convert label to node index # convert label to node index
label = d['eval'][2] label = d["eval"][2]
label_idx = nid2idx[label] label_idx = nid2idx[label]
labels.append(label_idx) labels.append(label_idx)
...@@ -71,19 +122,26 @@ def _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse ...@@ -71,19 +122,26 @@ def _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse
if e in reverse_edge: if e in reverse_edge:
g.add_edge(nid2idx[t], nid2idx[s]) g.add_edge(nid2idx[t], nid2idx[s])
edge_types.append(reverse_edge[e]) edge_types.append(reverse_edge[e])
g.edata['type'] = torch.tensor(edge_types, dtype=torch.long) g.edata["type"] = torch.tensor(edge_types, dtype=torch.long)
annotation = torch.zeros(len(node_ids), dtype=torch.long) annotation = torch.zeros(len(node_ids), dtype=torch.long)
annotation[nid2idx[d['eval'][0]]] = 1 annotation[nid2idx[d["eval"][0]]] = 1
g.ndata['annotation'] = annotation.unsqueeze(-1) g.ndata["annotation"] = annotation.unsqueeze(-1)
graphs.append(g) graphs.append(g)
batch_graph = dgl.batch(graphs) batch_graph = dgl.batch(graphs)
labels = torch.tensor(labels, dtype=torch.long) labels = torch.tensor(labels, dtype=torch.long)
return batch_graph, labels return batch_graph, labels
def _get_dataloader(data, shuffle): def _get_dataloader(data, shuffle):
return DataLoader(dataset=data, batch_size=batch_size, shuffle=shuffle, collate_fn=_collate_fn) return DataLoader(
dataset=data,
train_set, dev_set, test_sets = _convert_ns_dataset(train_size, node_dict, edge_dict, path, q_type) batch_size=batch_size,
shuffle=shuffle,
collate_fn=_collate_fn,
)
train_set, dev_set, test_sets = _convert_ns_dataset(
train_size, node_dict, edge_dict, path, q_type
)
train_dataloader = _get_dataloader(train_set, True) train_dataloader = _get_dataloader(train_set, True)
dev_dataloader = _get_dataloader(dev_set, False) dev_dataloader = _get_dataloader(dev_set, False)
test_dataloaders = [] test_dataloaders = []
...@@ -100,26 +158,36 @@ def _convert_ns_dataset(train_size, node_dict, edge_dict, path, q_type): ...@@ -100,26 +158,36 @@ def _convert_ns_dataset(train_size, node_dict, edge_dict, path, q_type):
def convert(file): def convert(file):
dataset = [] dataset = []
d = dict() d = dict()
with open(file, 'r') as f: with open(file, "r") as f:
for i, line in enumerate(f.readlines()): for i, line in enumerate(f.readlines()):
line = line.strip().split() line = line.strip().split()
if line[0] == '1' and len(d) > 0: if line[0] == "1" and len(d) > 0:
d = dict() d = dict()
if line[1] == 'eval': if line[1] == "eval":
# (src, edge, label) # (src, edge, label)
d['eval'] = (node_dict[line[2]], edge_dict[line[3]], node_dict[line[4]]) d["eval"] = (
if d['eval'][1] == q_type: node_dict[line[2]],
edge_dict[line[3]],
node_dict[line[4]],
)
if d["eval"][1] == q_type:
dataset.append(d) dataset.append(d)
if len(dataset) >= total_num: if len(dataset) >= total_num:
break break
else: else:
if 'edges' not in d: if "edges" not in d:
d['edges'] = [] d["edges"] = []
d['edges'].append((node_dict[line[1]], edge_dict[line[2]], node_dict[line[3]])) d["edges"].append(
(
node_dict[line[1]],
edge_dict[line[2]],
node_dict[line[3]],
)
)
return dataset return dataset
download_dir = get_download_dir() download_dir = get_download_dir()
filename = os.path.join(download_dir, 'babi_data', path, 'data.txt') filename = os.path.join(download_dir, "babi_data", path, "data.txt")
data = convert(filename) data = convert(filename)
assert len(data) == total_num assert len(data) == total_num
...@@ -128,18 +196,27 @@ def _convert_ns_dataset(train_size, node_dict, edge_dict, path, q_type): ...@@ -128,18 +196,27 @@ def _convert_ns_dataset(train_size, node_dict, edge_dict, path, q_type):
dev_set = data[950:1000] dev_set = data[950:1000]
test_sets = [] test_sets = []
for i in range(10): for i in range(10):
test = data[1000 * (i + 1): 1000 * (i + 2)] test = data[1000 * (i + 1) : 1000 * (i + 2)]
test_sets.append(test) test_sets.append(test)
return train_set, dev_set, test_sets return train_set, dev_set, test_sets
def _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_dict, reverse_edge, path): def _gc_dataloader(
train_size,
q_type,
batch_size,
node_dict,
edge_dict,
label_dict,
reverse_edge,
path,
):
def _collate_fn(batch): def _collate_fn(batch):
graphs = [] graphs = []
labels = [] labels = []
for d in batch: for d in batch:
edges = d['edges'] edges = d["edges"]
node_ids = [] node_ids = []
for s, e, t in edges: for s, e, t in edges:
...@@ -149,11 +226,11 @@ def _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_d ...@@ -149,11 +226,11 @@ def _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_d
node_ids.append(t) node_ids.append(t)
g = dgl.DGLGraph() g = dgl.DGLGraph()
g.add_nodes(len(node_ids)) g.add_nodes(len(node_ids))
g.ndata['node_id'] = torch.tensor(node_ids, dtype=torch.long) g.ndata["node_id"] = torch.tensor(node_ids, dtype=torch.long)
nid2idx = dict(zip(node_ids, list(range(len(node_ids))))) nid2idx = dict(zip(node_ids, list(range(len(node_ids)))))
labels.append(d['eval'][-1]) labels.append(d["eval"][-1])
edge_types = [] edge_types = []
for s, e, t in edges: for s, e, t in edges:
...@@ -162,20 +239,27 @@ def _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_d ...@@ -162,20 +239,27 @@ def _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_d
if e in reverse_edge: if e in reverse_edge:
g.add_edge(nid2idx[t], nid2idx[s]) g.add_edge(nid2idx[t], nid2idx[s])
edge_types.append(reverse_edge[e]) edge_types.append(reverse_edge[e])
g.edata['type'] = torch.tensor(edge_types, dtype=torch.long) g.edata["type"] = torch.tensor(edge_types, dtype=torch.long)
annotation = torch.zeros([len(node_ids), 2], dtype=torch.long) annotation = torch.zeros([len(node_ids), 2], dtype=torch.long)
annotation[nid2idx[d['eval'][0]]][0] = 1 annotation[nid2idx[d["eval"][0]]][0] = 1
annotation[nid2idx[d['eval'][2]]][1] = 1 annotation[nid2idx[d["eval"][2]]][1] = 1
g.ndata['annotation'] = annotation g.ndata["annotation"] = annotation
graphs.append(g) graphs.append(g)
batch_graph = dgl.batch(graphs) batch_graph = dgl.batch(graphs)
labels = torch.tensor(labels, dtype=torch.long) labels = torch.tensor(labels, dtype=torch.long)
return batch_graph, labels return batch_graph, labels
def _get_dataloader(data, shuffle): def _get_dataloader(data, shuffle):
return DataLoader(dataset=data, batch_size=batch_size, shuffle=shuffle, collate_fn=_collate_fn) return DataLoader(
dataset=data,
train_set, dev_set, test_sets = _convert_gc_dataset(train_size, node_dict, edge_dict, label_dict, path, q_type) batch_size=batch_size,
shuffle=shuffle,
collate_fn=_collate_fn,
)
train_set, dev_set, test_sets = _convert_gc_dataset(
train_size, node_dict, edge_dict, label_dict, path, q_type
)
train_dataloader = _get_dataloader(train_set, True) train_dataloader = _get_dataloader(train_set, True)
dev_dataloader = _get_dataloader(dev_set, False) dev_dataloader = _get_dataloader(dev_set, False)
test_dataloaders = [] test_dataloaders = []
...@@ -186,33 +270,46 @@ def _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_d ...@@ -186,33 +270,46 @@ def _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_d
return train_dataloader, dev_dataloader, test_dataloaders return train_dataloader, dev_dataloader, test_dataloaders
def _convert_gc_dataset(train_size, node_dict, edge_dict, label_dict, path, q_type): def _convert_gc_dataset(
train_size, node_dict, edge_dict, label_dict, path, q_type
):
total_num = 11000 total_num = 11000
def convert(file): def convert(file):
dataset = [] dataset = []
d = dict() d = dict()
with open(file, 'r') as f: with open(file, "r") as f:
for i, line in enumerate(f.readlines()): for i, line in enumerate(f.readlines()):
line = line.strip().split() line = line.strip().split()
if line[0] == '1' and len(d) > 0: if line[0] == "1" and len(d) > 0:
d = dict() d = dict()
if line[1] == 'eval': if line[1] == "eval":
# (src, edge, label) # (src, edge, label)
if 'eval' not in d: if "eval" not in d:
d['eval'] = (node_dict[line[2]], edge_dict[line[3]], node_dict[line[4]], label_dict[line[5]]) d["eval"] = (
if d['eval'][1] == q_type: node_dict[line[2]],
edge_dict[line[3]],
node_dict[line[4]],
label_dict[line[5]],
)
if d["eval"][1] == q_type:
dataset.append(d) dataset.append(d)
if len(dataset) >= total_num: if len(dataset) >= total_num:
break break
else: else:
if 'edges' not in d: if "edges" not in d:
d['edges'] = [] d["edges"] = []
d['edges'].append((node_dict[line[1]], edge_dict[line[2]], node_dict[line[3]])) d["edges"].append(
(
node_dict[line[1]],
edge_dict[line[2]],
node_dict[line[3]],
)
)
return dataset return dataset
download_dir = get_download_dir() download_dir = get_download_dir()
filename = os.path.join(download_dir, 'babi_data', path, 'data.txt') filename = os.path.join(download_dir, "babi_data", path, "data.txt")
data = convert(filename) data = convert(filename)
assert len(data) == total_num assert len(data) == total_num
...@@ -221,19 +318,27 @@ def _convert_gc_dataset(train_size, node_dict, edge_dict, label_dict, path, q_ty ...@@ -221,19 +318,27 @@ def _convert_gc_dataset(train_size, node_dict, edge_dict, label_dict, path, q_ty
dev_set = data[950:1000] dev_set = data[950:1000]
test_sets = [] test_sets = []
for i in range(10): for i in range(10):
test = data[1000 * (i + 1): 1000 * (i + 2)] test = data[1000 * (i + 1) : 1000 * (i + 2)]
test_sets.append(test) test_sets.append(test)
return train_set, dev_set, test_sets return train_set, dev_set, test_sets
def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, reverse_edge, path, max_seq_length): def _path_finding_dataloader(
train_size,
batch_size,
node_dict,
edge_dict,
reverse_edge,
path,
max_seq_length,
):
def _collate_fn(batch): def _collate_fn(batch):
graphs = [] graphs = []
ground_truths = [] ground_truths = []
seq_lengths = [] seq_lengths = []
for d in batch: for d in batch:
edges = d['edges'] edges = d["edges"]
node_ids = [] node_ids = []
for s, e, t in edges: for s, e, t in edges:
...@@ -243,12 +348,14 @@ def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, rever ...@@ -243,12 +348,14 @@ def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, rever
node_ids.append(t) node_ids.append(t)
g = dgl.DGLGraph() g = dgl.DGLGraph()
g.add_nodes(len(node_ids)) g.add_nodes(len(node_ids))
g.ndata['node_id'] = torch.tensor(node_ids, dtype=torch.long) g.ndata["node_id"] = torch.tensor(node_ids, dtype=torch.long)
nid2idx = dict(zip(node_ids, list(range(len(node_ids))))) nid2idx = dict(zip(node_ids, list(range(len(node_ids)))))
truth = d['seq_out'] + [edge_dict['<end>']] * (max_seq_length - len(d['seq_out'])) truth = d["seq_out"] + [edge_dict["<end>"]] * (
seq_len = len(d['seq_out']) max_seq_length - len(d["seq_out"])
)
seq_len = len(d["seq_out"])
ground_truths.append(truth) ground_truths.append(truth)
seq_lengths.append(seq_len) seq_lengths.append(seq_len)
...@@ -259,11 +366,11 @@ def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, rever ...@@ -259,11 +366,11 @@ def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, rever
if e in reverse_edge: if e in reverse_edge:
g.add_edge(nid2idx[t], nid2idx[s]) g.add_edge(nid2idx[t], nid2idx[s])
edge_types.append(reverse_edge[e]) edge_types.append(reverse_edge[e])
g.edata['type'] = torch.tensor(edge_types, dtype=torch.long) g.edata["type"] = torch.tensor(edge_types, dtype=torch.long)
annotation = torch.zeros([len(node_ids), 2], dtype=torch.long) annotation = torch.zeros([len(node_ids), 2], dtype=torch.long)
annotation[nid2idx[d['eval'][0]]][0] = 1 annotation[nid2idx[d["eval"][0]]][0] = 1
annotation[nid2idx[d['eval'][1]]][1] = 1 annotation[nid2idx[d["eval"][1]]][1] = 1
g.ndata['annotation'] = annotation g.ndata["annotation"] = annotation
graphs.append(g) graphs.append(g)
batch_graph = dgl.batch(graphs) batch_graph = dgl.batch(graphs)
ground_truths = torch.tensor(ground_truths, dtype=torch.long) ground_truths = torch.tensor(ground_truths, dtype=torch.long)
...@@ -271,9 +378,16 @@ def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, rever ...@@ -271,9 +378,16 @@ def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, rever
return batch_graph, ground_truths, seq_lengths return batch_graph, ground_truths, seq_lengths
def _get_dataloader(data, shuffle): def _get_dataloader(data, shuffle):
return DataLoader(dataset=data, batch_size=batch_size, shuffle=shuffle, collate_fn=_collate_fn) return DataLoader(
dataset=data,
train_set, dev_set, test_sets = _convert_path_finding(train_size, node_dict, edge_dict, path) batch_size=batch_size,
shuffle=shuffle,
collate_fn=_collate_fn,
)
train_set, dev_set, test_sets = _convert_path_finding(
train_size, node_dict, edge_dict, path
)
train_dataloader = _get_dataloader(train_set, True) train_dataloader = _get_dataloader(train_set, True)
dev_dataloader = _get_dataloader(dev_set, False) dev_dataloader = _get_dataloader(dev_set, False)
test_dataloaders = [] test_dataloaders = []
...@@ -290,29 +404,35 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path): ...@@ -290,29 +404,35 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path):
def convert(file): def convert(file):
dataset = [] dataset = []
d = dict() d = dict()
with open(file, 'r') as f: with open(file, "r") as f:
for line in f.readlines(): for line in f.readlines():
line = line.strip().split() line = line.strip().split()
if line[0] == '1' and len(d) > 0: if line[0] == "1" and len(d) > 0:
d = dict() d = dict()
if line[1] == 'eval': if line[1] == "eval":
# (src, edge, label) # (src, edge, label)
d['eval'] = (node_dict[line[3]], node_dict[line[4]]) d["eval"] = (node_dict[line[3]], node_dict[line[4]])
d['seq_out'] = [] d["seq_out"] = []
seq_out = line[5].split(',') seq_out = line[5].split(",")
for e in seq_out: for e in seq_out:
d['seq_out'].append(edge_dict[e]) d["seq_out"].append(edge_dict[e])
dataset.append(d) dataset.append(d)
if len(dataset) >= total_num: if len(dataset) >= total_num:
break break
else: else:
if 'edges' not in d: if "edges" not in d:
d['edges'] = [] d["edges"] = []
d['edges'].append((node_dict[line[1]], edge_dict[line[2]], node_dict[line[3]])) d["edges"].append(
(
node_dict[line[1]],
edge_dict[line[2]],
node_dict[line[3]],
)
)
return dataset return dataset
download_dir = get_download_dir() download_dir = get_download_dir()
filename = os.path.join(download_dir, 'babi_data', path, 'data.txt') filename = os.path.join(download_dir, "babi_data", path, "data.txt")
data = convert(filename) data = convert(filename)
assert len(data) == total_num assert len(data) == total_num
...@@ -321,7 +441,7 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path): ...@@ -321,7 +441,7 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path):
dev_set = data[950:1000] dev_set = data[950:1000]
test_sets = [] test_sets = []
for i in range(10): for i in range(10):
test = data[1000 * (i + 1): 1000 * (i + 2)] test = data[1000 * (i + 1) : 1000 * (i + 2)]
test_sets.append(test) test_sets.append(test)
return train_set, dev_set, test_sets return train_set, dev_set, test_sets
...@@ -329,11 +449,11 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path): ...@@ -329,11 +449,11 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path):
def _download_babi_data(): def _download_babi_data():
download_dir = get_download_dir() download_dir = get_download_dir()
zip_file_path = os.path.join(download_dir, 'babi_data.zip') zip_file_path = os.path.join(download_dir, "babi_data.zip")
data_url = _get_dgl_url('models/ggnn_babi_data.zip') data_url = _get_dgl_url("models/ggnn_babi_data.zip")
download(data_url, path=zip_file_path) download(data_url, path=zip_file_path)
extract_dir = os.path.join(download_dir, 'babi_data') extract_dir = os.path.join(download_dir, "babi_data")
if not os.path.exists(extract_dir): if not os.path.exists(extract_dir):
extract_archive(zip_file_path, extract_dir) extract_archive(zip_file_path, extract_dir)
\ No newline at end of file
""" """
Gated Graph Neural Network module for graph classification tasks Gated Graph Neural Network module for graph classification tasks
""" """
from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
import torch import torch
from torch import nn from torch import nn
from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
class GraphClsGGNN(nn.Module): class GraphClsGGNN(nn.Module):
def __init__(self, def __init__(self, annotation_size, out_feats, n_steps, n_etypes, num_cls):
annotation_size,
out_feats,
n_steps,
n_etypes,
num_cls):
super(GraphClsGGNN, self).__init__() super(GraphClsGGNN, self).__init__()
self.annotation_size = annotation_size self.annotation_size = annotation_size
self.out_feats = out_feats self.out_feats = out_feats
self.ggnn = GatedGraphConv(in_feats=out_feats, self.ggnn = GatedGraphConv(
out_feats=out_feats, in_feats=out_feats,
n_steps=n_steps, out_feats=out_feats,
n_etypes=n_etypes) n_steps=n_steps,
n_etypes=n_etypes,
)
pooling_gate_nn = nn.Linear(annotation_size + out_feats, 1) pooling_gate_nn = nn.Linear(annotation_size + out_feats, 1)
self.pooling = GlobalAttentionPooling(pooling_gate_nn) self.pooling = GlobalAttentionPooling(pooling_gate_nn)
...@@ -30,16 +28,18 @@ class GraphClsGGNN(nn.Module): ...@@ -30,16 +28,18 @@ class GraphClsGGNN(nn.Module):
self.loss_fn = nn.CrossEntropyLoss() self.loss_fn = nn.CrossEntropyLoss()
def forward(self, graph, labels=None): def forward(self, graph, labels=None):
etypes = graph.edata.pop('type') etypes = graph.edata.pop("type")
annotation = graph.ndata.pop('annotation').float() annotation = graph.ndata.pop("annotation").float()
assert annotation.size()[-1] == self.annotation_size assert annotation.size()[-1] == self.annotation_size
node_num = graph.number_of_nodes() node_num = graph.number_of_nodes()
zero_pad = torch.zeros([node_num, self.out_feats - self.annotation_size], zero_pad = torch.zeros(
dtype=torch.float, [node_num, self.out_feats - self.annotation_size],
device=annotation.device) dtype=torch.float,
device=annotation.device,
)
h1 = torch.cat([annotation, zero_pad], -1) h1 = torch.cat([annotation, zero_pad], -1)
out = self.ggnn(graph, h1, etypes) out = self.ggnn(graph, h1, etypes)
...@@ -54,4 +54,4 @@ class GraphClsGGNN(nn.Module): ...@@ -54,4 +54,4 @@ class GraphClsGGNN(nn.Module):
if labels is not None: if labels is not None:
loss = self.loss_fn(logits, labels) loss = self.loss_fn(logits, labels)
return loss, preds return loss, preds
return preds return preds
\ No newline at end of file
""" """
Gated Graph Neural Network module for node selection tasks Gated Graph Neural Network module for node selection tasks
""" """
from dgl.nn.pytorch import GatedGraphConv
import torch import torch
from torch import nn from torch import nn
import dgl import dgl
from dgl.nn.pytorch import GatedGraphConv
class NodeSelectionGGNN(nn.Module): class NodeSelectionGGNN(nn.Module):
def __init__(self, def __init__(self, annotation_size, out_feats, n_steps, n_etypes):
annotation_size,
out_feats,
n_steps,
n_etypes):
super(NodeSelectionGGNN, self).__init__() super(NodeSelectionGGNN, self).__init__()
self.annotation_size = annotation_size self.annotation_size = annotation_size
self.out_feats = out_feats self.out_feats = out_feats
self.ggnn = GatedGraphConv(in_feats=out_feats, self.ggnn = GatedGraphConv(
out_feats=out_feats, in_feats=out_feats,
n_steps=n_steps, out_feats=out_feats,
n_etypes=n_etypes) n_steps=n_steps,
n_etypes=n_etypes,
)
self.output_layer = nn.Linear(annotation_size + out_feats, 1) self.output_layer = nn.Linear(annotation_size + out_feats, 1)
self.loss_fn = nn.CrossEntropyLoss() self.loss_fn = nn.CrossEntropyLoss()
def forward(self, graph, labels=None): def forward(self, graph, labels=None):
etypes = graph.edata.pop('type') etypes = graph.edata.pop("type")
annotation = graph.ndata.pop('annotation').float() annotation = graph.ndata.pop("annotation").float()
assert annotation.size()[-1] == self.annotation_size assert annotation.size()[-1] == self.annotation_size
node_num = graph.number_of_nodes() node_num = graph.number_of_nodes()
zero_pad = torch.zeros([node_num, self.out_feats - self.annotation_size], zero_pad = torch.zeros(
dtype=torch.float, [node_num, self.out_feats - self.annotation_size],
device=annotation.device) dtype=torch.float,
device=annotation.device,
)
h1 = torch.cat([annotation, zero_pad], -1) h1 = torch.cat([annotation, zero_pad], -1)
out = self.ggnn(graph, h1, etypes) out = self.ggnn(graph, h1, etypes)
all_logits = self.output_layer(torch.cat([out, annotation], -1)).squeeze(-1) all_logits = self.output_layer(
graph.ndata['logits'] = all_logits torch.cat([out, annotation], -1)
).squeeze(-1)
graph.ndata["logits"] = all_logits
batch_g = dgl.unbatch(graph) batch_g = dgl.unbatch(graph)
...@@ -50,7 +53,7 @@ class NodeSelectionGGNN(nn.Module): ...@@ -50,7 +53,7 @@ class NodeSelectionGGNN(nn.Module):
if labels is not None: if labels is not None:
loss = 0.0 loss = 0.0
for i, g in enumerate(batch_g): for i, g in enumerate(batch_g):
logits = g.ndata['logits'] logits = g.ndata["logits"]
preds.append(torch.argmax(logits)) preds.append(torch.argmax(logits))
if labels is not None: if labels is not None:
logits = logits.unsqueeze(0) logits = logits.unsqueeze(0)
...@@ -60,4 +63,4 @@ class NodeSelectionGGNN(nn.Module): ...@@ -60,4 +63,4 @@ class NodeSelectionGGNN(nn.Module):
if labels is not None: if labels is not None:
loss /= float(len(batch_g)) loss /= float(len(batch_g))
return loss, preds return loss, preds
return preds return preds
\ No newline at end of file
...@@ -2,42 +2,49 @@ ...@@ -2,42 +2,49 @@
Gated Graph Sequence Neural Network for sequence outputs Gated Graph Sequence Neural Network for sequence outputs
""" """
from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
import torch import torch
from torch import nn
import torch.nn.functional as F import torch.nn.functional as F
from torch import nn
from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
class GGSNN(nn.Module): class GGSNN(nn.Module):
def __init__(self, def __init__(
annotation_size, self,
out_feats, annotation_size,
n_steps, out_feats,
n_etypes, n_steps,
max_seq_length, n_etypes,
num_cls): max_seq_length,
num_cls,
):
super(GGSNN, self).__init__() super(GGSNN, self).__init__()
self.annotation_size = annotation_size self.annotation_size = annotation_size
self.out_feats = out_feats self.out_feats = out_feats
self.max_seq_length = max_seq_length self.max_seq_length = max_seq_length
self.ggnn = GatedGraphConv(in_feats=out_feats, self.ggnn = GatedGraphConv(
out_feats=out_feats, in_feats=out_feats,
n_steps=n_steps, out_feats=out_feats,
n_etypes=n_etypes) n_steps=n_steps,
n_etypes=n_etypes,
)
self.annotation_out_layer = nn.Linear(annotation_size + out_feats, annotation_size) self.annotation_out_layer = nn.Linear(
annotation_size + out_feats, annotation_size
)
pooling_gate_nn = nn.Linear(annotation_size + out_feats, 1) pooling_gate_nn = nn.Linear(annotation_size + out_feats, 1)
self.pooling = GlobalAttentionPooling(pooling_gate_nn) self.pooling = GlobalAttentionPooling(pooling_gate_nn)
self.output_layer = nn.Linear(annotation_size + out_feats, num_cls) self.output_layer = nn.Linear(annotation_size + out_feats, num_cls)
self.loss_fn = nn.CrossEntropyLoss(reduction='none') self.loss_fn = nn.CrossEntropyLoss(reduction="none")
def forward(self, graph, seq_lengths, ground_truth=None): def forward(self, graph, seq_lengths, ground_truth=None):
etypes = graph.edata.pop('type') etypes = graph.edata.pop("type")
annotation = graph.ndata.pop('annotation').float() annotation = graph.ndata.pop("annotation").float()
assert annotation.size()[-1] == self.annotation_size assert annotation.size()[-1] == self.annotation_size
...@@ -45,9 +52,11 @@ class GGSNN(nn.Module): ...@@ -45,9 +52,11 @@ class GGSNN(nn.Module):
all_logits = [] all_logits = []
for _ in range(self.max_seq_length): for _ in range(self.max_seq_length):
zero_pad = torch.zeros([node_num, self.out_feats - self.annotation_size], zero_pad = torch.zeros(
dtype=torch.float, [node_num, self.out_feats - self.annotation_size],
device=annotation.device) dtype=torch.float,
device=annotation.device,
)
h1 = torch.cat([annotation.detach(), zero_pad], -1) h1 = torch.cat([annotation.detach(), zero_pad], -1)
out = self.ggnn(graph, h1, etypes) out = self.ggnn(graph, h1, etypes)
...@@ -71,14 +80,18 @@ def sequence_loss(logits, ground_truth, seq_length=None): ...@@ -71,14 +80,18 @@ def sequence_loss(logits, ground_truth, seq_length=None):
def sequence_mask(length): def sequence_mask(length):
max_length = logits.size(1) max_length = logits.size(1)
batch_size = logits.size(0) batch_size = logits.size(0)
range_tensor = torch.arange(0, max_length, dtype=seq_length.dtype, device=seq_length.device) range_tensor = torch.arange(
range_tensor = torch.stack([range_tensor]*batch_size, 0) 0, max_length, dtype=seq_length.dtype, device=seq_length.device
)
range_tensor = torch.stack([range_tensor] * batch_size, 0)
expanded_length = torch.stack([length]*max_length, -1) expanded_length = torch.stack([length] * max_length, -1)
mask = (range_tensor < expanded_length).float() mask = (range_tensor < expanded_length).float()
return mask return mask
loss = nn.CrossEntropyLoss(reduction='none')(logits.permute((0, 2, 1)), ground_truth) loss = nn.CrossEntropyLoss(reduction="none")(
logits.permute((0, 2, 1)), ground_truth
)
if seq_length is None: if seq_length is None:
loss = loss.mean() loss = loss.mean()
...@@ -86,4 +99,4 @@ def sequence_loss(logits, ground_truth, seq_length=None): ...@@ -86,4 +99,4 @@ def sequence_loss(logits, ground_truth, seq_length=None):
mask = sequence_mask(seq_length) mask = sequence_mask(seq_length)
loss = (loss * mask).sum(-1) / seq_length.float() loss = (loss * mask).sum(-1) / seq_length.float()
loss = loss.mean() loss = loss.mean()
return loss return loss
\ No newline at end of file
...@@ -3,34 +3,38 @@ Training and testing for graph classification tasks in bAbI ...@@ -3,34 +3,38 @@ Training and testing for graph classification tasks in bAbI
""" """
import argparse import argparse
import numpy as np
import torch
from data_utils import get_babi_dataloaders from data_utils import get_babi_dataloaders
from ggnn_gc import GraphClsGGNN from ggnn_gc import GraphClsGGNN
from torch.optim import Adam from torch.optim import Adam
import torch
import numpy as np
def main(args): def main(args):
out_feats = {18: 3} out_feats = {18: 3}
n_etypes = {18: 2} n_etypes = {18: 2}
train_dataloader, dev_dataloader, test_dataloaders = \ train_dataloader, dev_dataloader, test_dataloaders = get_babi_dataloaders(
get_babi_dataloaders(batch_size=args.batch_size, batch_size=args.batch_size,
train_size=args.train_num, train_size=args.train_num,
task_id=args.task_id, task_id=args.task_id,
q_type=args.question_id) q_type=args.question_id,
)
model = GraphClsGGNN(annotation_size=2,
out_feats=out_feats[args.task_id], model = GraphClsGGNN(
n_steps=5, annotation_size=2,
n_etypes=n_etypes[args.task_id], out_feats=out_feats[args.task_id],
num_cls=2) n_steps=5,
n_etypes=n_etypes[args.task_id],
num_cls=2,
)
opt = Adam(model.parameters(), lr=args.lr) opt = Adam(model.parameters(), lr=args.lr)
print(f'Task {args.task_id}, question_id {args.question_id}') print(f"Task {args.task_id}, question_id {args.question_id}")
print(f'Training set size: {len(train_dataloader.dataset)}') print(f"Training set size: {len(train_dataloader.dataset)}")
print(f'Dev set size: {len(dev_dataloader.dataset)}') print(f"Dev set size: {len(dev_dataloader.dataset)}")
# training and dev stage # training and dev stage
for epoch in range(args.epochs): for epoch in range(args.epochs):
...@@ -42,7 +46,7 @@ def main(args): ...@@ -42,7 +46,7 @@ def main(args):
loss.backward() loss.backward()
opt.step() opt.step()
if epoch % 20 == 0: if epoch % 20 == 0:
print(f'Epoch {epoch}, batch {i} loss: {loss.data}') print(f"Epoch {epoch}, batch {i} loss: {loss.data}")
if epoch % 20 != 0: if epoch % 20 != 0:
continue continue
...@@ -62,7 +66,7 @@ def main(args): ...@@ -62,7 +66,7 @@ def main(args):
# test stage # test stage
for i, dataloader in enumerate(test_dataloaders): for i, dataloader in enumerate(test_dataloaders):
print(f'Test set {i} size: {len(dataloader.dataset)}') print(f"Test set {i} size: {len(dataloader.dataset)}")
test_acc_list = [] test_acc_list = []
for dataloader in test_dataloaders: for dataloader in test_dataloaders:
...@@ -83,24 +87,30 @@ def main(args): ...@@ -83,24 +87,30 @@ def main(args):
test_acc_mean = np.mean(test_acc_list) test_acc_mean = np.mean(test_acc_list)
test_acc_std = np.std(test_acc_list) test_acc_std = np.std(test_acc_list)
print(f'Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}') print(
f"Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}"
)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Gated Graph Neural Networks for graph classification tasks in bAbI')
parser.add_argument('--task_id', type=int, default=18, if __name__ == "__main__":
help='task id from 1 to 20') parser = argparse.ArgumentParser(
parser.add_argument('--question_id', type=int, default=0, description="Gated Graph Neural Networks for graph classification tasks in bAbI"
help='question id for each task') )
parser.add_argument('--train_num', type=int, default=950, parser.add_argument(
help='Number of training examples') "--task_id", type=int, default=18, help="task id from 1 to 20"
parser.add_argument('--batch_size', type=int, default=50, )
help='batch size') parser.add_argument(
parser.add_argument('--lr', type=float, default=1e-3, "--question_id", type=int, default=0, help="question id for each task"
help='learning rate') )
parser.add_argument('--epochs', type=int, default=200, parser.add_argument(
help='number of training epochs') "--train_num", type=int, default=950, help="Number of training examples"
)
parser.add_argument("--batch_size", type=int, default=50, help="batch size")
parser.add_argument("--lr", type=float, default=1e-3, help="learning rate")
parser.add_argument(
"--epochs", type=int, default=200, help="number of training epochs"
)
args = parser.parse_args() args = parser.parse_args()
main(args) main(args)
\ No newline at end of file
...@@ -3,34 +3,38 @@ Training and testing for node selection tasks in bAbI ...@@ -3,34 +3,38 @@ Training and testing for node selection tasks in bAbI
""" """
import argparse import argparse
import time
import numpy as np
import torch
from data_utils import get_babi_dataloaders from data_utils import get_babi_dataloaders
from ggnn_ns import NodeSelectionGGNN from ggnn_ns import NodeSelectionGGNN
from torch.optim import Adam from torch.optim import Adam
import torch
import numpy as np
import time
def main(args): def main(args):
out_feats = {4: 4, 15: 5, 16: 6} out_feats = {4: 4, 15: 5, 16: 6}
n_etypes = {4: 4, 15: 2, 16: 2} n_etypes = {4: 4, 15: 2, 16: 2}
train_dataloader, dev_dataloader, test_dataloaders = \ train_dataloader, dev_dataloader, test_dataloaders = get_babi_dataloaders(
get_babi_dataloaders(batch_size=args.batch_size, batch_size=args.batch_size,
train_size=args.train_num, train_size=args.train_num,
task_id=args.task_id, task_id=args.task_id,
q_type=args.question_id) q_type=args.question_id,
)
model = NodeSelectionGGNN(annotation_size=1,
out_feats=out_feats[args.task_id], model = NodeSelectionGGNN(
n_steps=5, annotation_size=1,
n_etypes=n_etypes[args.task_id]) out_feats=out_feats[args.task_id],
n_steps=5,
n_etypes=n_etypes[args.task_id],
)
opt = Adam(model.parameters(), lr=args.lr) opt = Adam(model.parameters(), lr=args.lr)
print(f'Task {args.task_id}, question_id {args.question_id}') print(f"Task {args.task_id}, question_id {args.question_id}")
print(f'Training set size: {len(train_dataloader.dataset)}') print(f"Training set size: {len(train_dataloader.dataset)}")
print(f'Dev set size: {len(dev_dataloader.dataset)}') print(f"Dev set size: {len(dev_dataloader.dataset)}")
# training and dev stage # training and dev stage
for epoch in range(args.epochs): for epoch in range(args.epochs):
...@@ -41,7 +45,7 @@ def main(args): ...@@ -41,7 +45,7 @@ def main(args):
opt.zero_grad() opt.zero_grad()
loss.backward() loss.backward()
opt.step() opt.step()
print(f'Epoch {epoch}, batch {i} loss: {loss.data}') print(f"Epoch {epoch}, batch {i} loss: {loss.data}")
dev_preds = [] dev_preds = []
dev_labels = [] dev_labels = []
...@@ -49,7 +53,9 @@ def main(args): ...@@ -49,7 +53,9 @@ def main(args):
for g, labels in dev_dataloader: for g, labels in dev_dataloader:
with torch.no_grad(): with torch.no_grad():
preds = model(g) preds = model(g)
preds = torch.tensor(preds, dtype=torch.long).data.numpy().tolist() preds = (
torch.tensor(preds, dtype=torch.long).data.numpy().tolist()
)
labels = labels.data.numpy().tolist() labels = labels.data.numpy().tolist()
dev_preds += preds dev_preds += preds
dev_labels += labels dev_labels += labels
...@@ -59,7 +65,7 @@ def main(args): ...@@ -59,7 +65,7 @@ def main(args):
# test stage # test stage
for i, dataloader in enumerate(test_dataloaders): for i, dataloader in enumerate(test_dataloaders):
print(f'Test set {i} size: {len(dataloader.dataset)}') print(f"Test set {i} size: {len(dataloader.dataset)}")
test_acc_list = [] test_acc_list = []
for dataloader in test_dataloaders: for dataloader in test_dataloaders:
...@@ -69,7 +75,9 @@ def main(args): ...@@ -69,7 +75,9 @@ def main(args):
for g, labels in dataloader: for g, labels in dataloader:
with torch.no_grad(): with torch.no_grad():
preds = model(g) preds = model(g)
preds = torch.tensor(preds, dtype=torch.long).data.numpy().tolist() preds = (
torch.tensor(preds, dtype=torch.long).data.numpy().tolist()
)
labels = labels.data.numpy().tolist() labels = labels.data.numpy().tolist()
test_preds += preds test_preds += preds
test_labels += labels test_labels += labels
...@@ -80,24 +88,30 @@ def main(args): ...@@ -80,24 +88,30 @@ def main(args):
test_acc_mean = np.mean(test_acc_list) test_acc_mean = np.mean(test_acc_list)
test_acc_std = np.std(test_acc_list) test_acc_std = np.std(test_acc_list)
print(f'Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}') print(
f"Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}"
)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Gated Graph Neural Networks for node selection tasks in bAbI')
parser.add_argument('--task_id', type=int, default=16, if __name__ == "__main__":
help='task id from 1 to 20') parser = argparse.ArgumentParser(
parser.add_argument('--question_id', type=int, default=1, description="Gated Graph Neural Networks for node selection tasks in bAbI"
help='question id for each task') )
parser.add_argument('--train_num', type=int, default=50, parser.add_argument(
help='Number of training examples') "--task_id", type=int, default=16, help="task id from 1 to 20"
parser.add_argument('--batch_size', type=int, default=10, )
help='batch size') parser.add_argument(
parser.add_argument('--lr', type=float, default=1e-3, "--question_id", type=int, default=1, help="question id for each task"
help='learning rate') )
parser.add_argument('--epochs', type=int, default=100, parser.add_argument(
help='number of training epochs') "--train_num", type=int, default=50, help="Number of training examples"
)
parser.add_argument("--batch_size", type=int, default=10, help="batch size")
parser.add_argument("--lr", type=float, default=1e-3, help="learning rate")
parser.add_argument(
"--epochs", type=int, default=100, help="number of training epochs"
)
args = parser.parse_args() args = parser.parse_args()
main(args) main(args)
\ No newline at end of file
...@@ -4,35 +4,39 @@ Here we take task 19 'Path Finding' as an example ...@@ -4,35 +4,39 @@ Here we take task 19 'Path Finding' as an example
""" """
import argparse import argparse
import numpy as np
import torch
from data_utils import get_babi_dataloaders from data_utils import get_babi_dataloaders
from ggsnn import GGSNN from ggsnn import GGSNN
from torch.optim import Adam from torch.optim import Adam
import torch
import numpy as np
def main(args): def main(args):
out_feats = {19: 6} out_feats = {19: 6}
n_etypes = {19: 4} n_etypes = {19: 4}
train_dataloader, dev_dataloader, test_dataloaders = \ train_dataloader, dev_dataloader, test_dataloaders = get_babi_dataloaders(
get_babi_dataloaders(batch_size=args.batch_size, batch_size=args.batch_size,
train_size=args.train_num, train_size=args.train_num,
task_id=args.task_id, task_id=args.task_id,
q_type=-1) q_type=-1,
)
model = GGSNN(annotation_size=2,
out_feats=out_feats[args.task_id], model = GGSNN(
n_steps=5, annotation_size=2,
n_etypes=n_etypes[args.task_id], out_feats=out_feats[args.task_id],
max_seq_length=2, n_steps=5,
num_cls=5) n_etypes=n_etypes[args.task_id],
max_seq_length=2,
num_cls=5,
)
opt = Adam(model.parameters(), lr=args.lr) opt = Adam(model.parameters(), lr=args.lr)
print(f'Task {args.task_id}') print(f"Task {args.task_id}")
print(f'Training set size: {len(train_dataloader.dataset)}') print(f"Training set size: {len(train_dataloader.dataset)}")
print(f'Dev set size: {len(dev_dataloader.dataset)}') print(f"Dev set size: {len(dev_dataloader.dataset)}")
# training and dev stage # training and dev stage
for epoch in range(args.epochs): for epoch in range(args.epochs):
...@@ -44,7 +48,7 @@ def main(args): ...@@ -44,7 +48,7 @@ def main(args):
loss.backward() loss.backward()
opt.step() opt.step()
if epoch % 20 == 0: if epoch % 20 == 0:
print(f'Epoch {epoch}, batch {i} loss: {loss.data}') print(f"Epoch {epoch}, batch {i} loss: {loss.data}")
if epoch % 20 != 0: if epoch % 20 != 0:
continue continue
...@@ -65,7 +69,7 @@ def main(args): ...@@ -65,7 +69,7 @@ def main(args):
# test stage # test stage
for i, dataloader in enumerate(test_dataloaders): for i, dataloader in enumerate(test_dataloaders):
print(f'Test set {i} size: {len(dataloader.dataset)}') print(f"Test set {i} size: {len(dataloader.dataset)}")
test_acc_list = [] test_acc_list = []
for dataloader in test_dataloaders: for dataloader in test_dataloaders:
...@@ -87,23 +91,28 @@ def main(args): ...@@ -87,23 +91,28 @@ def main(args):
test_acc_mean = np.mean(test_acc_list) test_acc_mean = np.mean(test_acc_list)
test_acc_std = np.std(test_acc_list) test_acc_std = np.std(test_acc_list)
print(f'Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}') print(
f"Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}"
)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Gated Graph Sequence Neural Networks for sequential output tasks in '
'bAbI') if __name__ == "__main__":
parser.add_argument('--task_id', type=int, default=19, parser = argparse.ArgumentParser(
help='task id from 1 to 20') description="Gated Graph Sequence Neural Networks for sequential output tasks in "
parser.add_argument('--train_num', type=int, default=250, "bAbI"
help='Number of training examples') )
parser.add_argument('--batch_size', type=int, default=10, parser.add_argument(
help='batch size') "--task_id", type=int, default=19, help="task id from 1 to 20"
parser.add_argument('--lr', type=float, default=1e-3, )
help='learning rate') parser.add_argument(
parser.add_argument('--epochs', type=int, default=200, "--train_num", type=int, default=250, help="Number of training examples"
help='number of training epochs') )
parser.add_argument("--batch_size", type=int, default=10, help="batch size")
parser.add_argument("--lr", type=float, default=1e-3, help="learning rate")
parser.add_argument(
"--epochs", type=int, default=200, help="number of training epochs"
)
args = parser.parse_args() args = parser.parse_args()
main(args) main(args)
\ No newline at end of file
import argparse
import numpy as np import numpy as np
import torch import torch
import torch.nn as nn import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F import torch.nn.functional as F
from torch.utils.data.sampler import SubsetRandomSampler import torch.optim as optim
from sklearn.model_selection import StratifiedKFold from sklearn.model_selection import StratifiedKFold
from torch.utils.data.sampler import SubsetRandomSampler
from dgl.data import GINDataset from dgl.data import GINDataset
from dgl.dataloading import GraphDataLoader from dgl.dataloading import GraphDataLoader
from dgl.nn.pytorch.conv import GINConv from dgl.nn.pytorch.conv import GINConv
from dgl.nn.pytorch.glob import SumPooling from dgl.nn.pytorch.glob import SumPooling
import argparse
class MLP(nn.Module): class MLP(nn.Module):
"""Construct two-layer MLP-type aggreator for GIN model""" """Construct two-layer MLP-type aggreator for GIN model"""
def __init__(self, input_dim, hidden_dim, output_dim): def __init__(self, input_dim, hidden_dim, output_dim):
super().__init__() super().__init__()
self.linears = nn.ModuleList() self.linears = nn.ModuleList()
# two-layer MLP # two-layer MLP
self.linears.append(nn.Linear(input_dim, hidden_dim, bias=False)) self.linears.append(nn.Linear(input_dim, hidden_dim, bias=False))
self.linears.append(nn.Linear(hidden_dim, output_dim, bias=False)) self.linears.append(nn.Linear(hidden_dim, output_dim, bias=False))
self.batch_norm = nn.BatchNorm1d((hidden_dim)) self.batch_norm = nn.BatchNorm1d((hidden_dim))
...@@ -25,7 +29,8 @@ class MLP(nn.Module): ...@@ -25,7 +29,8 @@ class MLP(nn.Module):
h = x h = x
h = F.relu(self.batch_norm(self.linears[0](h))) h = F.relu(self.batch_norm(self.linears[0](h)))
return self.linears[1](h) return self.linears[1](h)
class GIN(nn.Module): class GIN(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim): def __init__(self, input_dim, hidden_dim, output_dim):
super().__init__() super().__init__()
...@@ -33,12 +38,14 @@ class GIN(nn.Module): ...@@ -33,12 +38,14 @@ class GIN(nn.Module):
self.batch_norms = nn.ModuleList() self.batch_norms = nn.ModuleList()
num_layers = 5 num_layers = 5
# five-layer GCN with two-layer MLP aggregator and sum-neighbor-pooling scheme # five-layer GCN with two-layer MLP aggregator and sum-neighbor-pooling scheme
for layer in range(num_layers - 1): # excluding the input layer for layer in range(num_layers - 1): # excluding the input layer
if layer == 0: if layer == 0:
mlp = MLP(input_dim, hidden_dim, hidden_dim) mlp = MLP(input_dim, hidden_dim, hidden_dim)
else: else:
mlp = MLP(hidden_dim, hidden_dim, hidden_dim) mlp = MLP(hidden_dim, hidden_dim, hidden_dim)
self.ginlayers.append(GINConv(mlp, learn_eps=False)) # set to True if learning epsilon self.ginlayers.append(
GINConv(mlp, learn_eps=False)
) # set to True if learning epsilon
self.batch_norms.append(nn.BatchNorm1d(hidden_dim)) self.batch_norms.append(nn.BatchNorm1d(hidden_dim))
# linear functions for graph sum poolings of output of each layer # linear functions for graph sum poolings of output of each layer
self.linear_prediction = nn.ModuleList() self.linear_prediction = nn.ModuleList()
...@@ -48,7 +55,9 @@ class GIN(nn.Module): ...@@ -48,7 +55,9 @@ class GIN(nn.Module):
else: else:
self.linear_prediction.append(nn.Linear(hidden_dim, output_dim)) self.linear_prediction.append(nn.Linear(hidden_dim, output_dim))
self.drop = nn.Dropout(0.5) self.drop = nn.Dropout(0.5)
self.pool = SumPooling() # change to mean readout (AvgPooling) on social network datasets self.pool = (
SumPooling()
) # change to mean readout (AvgPooling) on social network datasets
def forward(self, g, h): def forward(self, g, h):
# list of hidden representation at each layer (including the input layer) # list of hidden representation at each layer (including the input layer)
...@@ -64,7 +73,8 @@ class GIN(nn.Module): ...@@ -64,7 +73,8 @@ class GIN(nn.Module):
pooled_h = self.pool(g, h) pooled_h = self.pool(g, h)
score_over_layer += self.drop(self.linear_prediction[i](pooled_h)) score_over_layer += self.drop(self.linear_prediction[i](pooled_h))
return score_over_layer return score_over_layer
def split_fold10(labels, fold_idx=0): def split_fold10(labels, fold_idx=0):
skf = StratifiedKFold(n_splits=10, shuffle=True, random_state=0) skf = StratifiedKFold(n_splits=10, shuffle=True, random_state=0)
idx_list = [] idx_list = []
...@@ -73,6 +83,7 @@ def split_fold10(labels, fold_idx=0): ...@@ -73,6 +83,7 @@ def split_fold10(labels, fold_idx=0):
train_idx, valid_idx = idx_list[fold_idx] train_idx, valid_idx = idx_list[fold_idx]
return train_idx, valid_idx return train_idx, valid_idx
def evaluate(dataloader, device, model): def evaluate(dataloader, device, model):
model.eval() model.eval()
total = 0 total = 0
...@@ -80,7 +91,7 @@ def evaluate(dataloader, device, model): ...@@ -80,7 +91,7 @@ def evaluate(dataloader, device, model):
for batched_graph, labels in dataloader: for batched_graph, labels in dataloader:
batched_graph = batched_graph.to(device) batched_graph = batched_graph.to(device)
labels = labels.to(device) labels = labels.to(device)
feat = batched_graph.ndata.pop('attr') feat = batched_graph.ndata.pop("attr")
total += len(labels) total += len(labels)
logits = model(batched_graph, feat) logits = model(batched_graph, feat)
_, predicted = torch.max(logits, 1) _, predicted = torch.max(logits, 1)
...@@ -88,20 +99,21 @@ def evaluate(dataloader, device, model): ...@@ -88,20 +99,21 @@ def evaluate(dataloader, device, model):
acc = 1.0 * total_correct / total acc = 1.0 * total_correct / total
return acc return acc
def train(train_loader, val_loader, device, model): def train(train_loader, val_loader, device, model):
# loss function, optimizer and scheduler # loss function, optimizer and scheduler
loss_fcn = nn.CrossEntropyLoss() loss_fcn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01) optimizer = optim.Adam(model.parameters(), lr=0.01)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5)
# training loop # training loop
for epoch in range(350): for epoch in range(350):
model.train() model.train()
total_loss = 0 total_loss = 0
for batch, (batched_graph, labels) in enumerate(train_loader): for batch, (batched_graph, labels) in enumerate(train_loader):
batched_graph = batched_graph.to(device) batched_graph = batched_graph.to(device)
labels = labels.to(device) labels = labels.to(device)
feat = batched_graph.ndata.pop('attr') feat = batched_graph.ndata.pop("attr")
logits = model(batched_graph, feat) logits = model(batched_graph, feat)
loss = loss_fcn(logits, labels) loss = loss_fcn(logits, labels)
optimizer.zero_grad() optimizer.zero_grad()
...@@ -111,34 +123,52 @@ def train(train_loader, val_loader, device, model): ...@@ -111,34 +123,52 @@ def train(train_loader, val_loader, device, model):
scheduler.step() scheduler.step()
train_acc = evaluate(train_loader, device, model) train_acc = evaluate(train_loader, device, model)
valid_acc = evaluate(val_loader, device, model) valid_acc = evaluate(val_loader, device, model)
print("Epoch {:05d} | Loss {:.4f} | Train Acc. {:.4f} | Validation Acc. {:.4f} " print(
. format(epoch, total_loss / (batch + 1), train_acc, valid_acc)) "Epoch {:05d} | Loss {:.4f} | Train Acc. {:.4f} | Validation Acc. {:.4f} ".format(
epoch, total_loss / (batch + 1), train_acc, valid_acc
)
)
if __name__ == '__main__': if __name__ == "__main__":
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default="MUTAG", parser.add_argument(
choices=['MUTAG', 'PTC', 'NCI1', 'PROTEINS'], "--dataset",
help='name of dataset (default: MUTAG)') type=str,
default="MUTAG",
choices=["MUTAG", "PTC", "NCI1", "PROTEINS"],
help="name of dataset (default: MUTAG)",
)
args = parser.parse_args() args = parser.parse_args()
print(f'Training with DGL built-in GINConv module with a fixed epsilon = 0') print(f"Training with DGL built-in GINConv module with a fixed epsilon = 0")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# load and split dataset # load and split dataset
dataset = GINDataset(args.dataset, self_loop=True, degree_as_nlabel=False) # add self_loop and disable one-hot encoding for input features dataset = GINDataset(
args.dataset, self_loop=True, degree_as_nlabel=False
) # add self_loop and disable one-hot encoding for input features
labels = [l for _, l in dataset] labels = [l for _, l in dataset]
train_idx, val_idx = split_fold10(labels) train_idx, val_idx = split_fold10(labels)
# create dataloader # create dataloader
train_loader = GraphDataLoader(dataset, sampler=SubsetRandomSampler(train_idx), train_loader = GraphDataLoader(
batch_size=128, pin_memory=torch.cuda.is_available()) dataset,
val_loader = GraphDataLoader(dataset, sampler=SubsetRandomSampler(val_idx), sampler=SubsetRandomSampler(train_idx),
batch_size=128, pin_memory=torch.cuda.is_available()) batch_size=128,
pin_memory=torch.cuda.is_available(),
)
val_loader = GraphDataLoader(
dataset,
sampler=SubsetRandomSampler(val_idx),
batch_size=128,
pin_memory=torch.cuda.is_available(),
)
# create GIN model # create GIN model
in_size = dataset.dim_nfeats in_size = dataset.dim_nfeats
out_size = dataset.gclasses out_size = dataset.gclasses
model = GIN(in_size, 16, out_size).to(device) model = GIN(in_size, 16, out_size).to(device)
# model training/validating # model training/validating
print('Training...') print("Training...")
train(train_loader, val_loader, device, model) train(train_loader, val_loader, device, model)
# Data augmentation on graphs via edge dropping and feature masking # Data augmentation on graphs via edge dropping and feature masking
import torch as th
import numpy as np import numpy as np
import torch as th
import dgl import dgl
def aug(graph, x, feat_drop_rate, edge_mask_rate): def aug(graph, x, feat_drop_rate, edge_mask_rate):
n_node = graph.num_nodes() n_node = graph.num_nodes()
...@@ -21,19 +23,22 @@ def aug(graph, x, feat_drop_rate, edge_mask_rate): ...@@ -21,19 +23,22 @@ def aug(graph, x, feat_drop_rate, edge_mask_rate):
return ng, feat return ng, feat
def drop_feature(x, drop_prob): def drop_feature(x, drop_prob):
drop_mask = th.empty((x.size(1),), drop_mask = (
dtype=th.float32, th.empty((x.size(1),), dtype=th.float32, device=x.device).uniform_(0, 1)
device=x.device).uniform_(0, 1) < drop_prob < drop_prob
)
x = x.clone() x = x.clone()
x[:, drop_mask] = 0 x[:, drop_mask] = 0
return x return x
def mask_edge(graph, mask_prob): def mask_edge(graph, mask_prob):
E = graph.num_edges() E = graph.num_edges()
mask_rates = th.FloatTensor(np.ones(E) * mask_prob) mask_rates = th.FloatTensor(np.ones(E) * mask_prob)
masks = th.bernoulli(1 - mask_rates) masks = th.bernoulli(1 - mask_rates)
mask_idx = masks.nonzero().squeeze(1) mask_idx = masks.nonzero().squeeze(1)
return mask_idx return mask_idx
\ No newline at end of file
from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
def load(name): def load(name):
if name == 'cora': if name == "cora":
dataset = CoraGraphDataset() dataset = CoraGraphDataset()
elif name == 'citeseer': elif name == "citeseer":
dataset = CiteseerGraphDataset() dataset = CiteseerGraphDataset()
elif name == 'pubmed': elif name == "pubmed":
dataset = PubmedGraphDataset() dataset = PubmedGraphDataset()
graph = dataset[0] graph = dataset[0]
train_mask = graph.ndata.pop('train_mask') train_mask = graph.ndata.pop("train_mask")
test_mask = graph.ndata.pop('test_mask') test_mask = graph.ndata.pop("test_mask")
feat = graph.ndata.pop('feat') feat = graph.ndata.pop("feat")
labels = graph.ndata.pop('label') labels = graph.ndata.pop("label")
return graph, feat, labels, train_mask, test_mask return graph, feat, labels, train_mask, test_mask
\ No newline at end of file
''' """
Code adapted from https://github.com/CRIPAC-DIG/GRACE Code adapted from https://github.com/CRIPAC-DIG/GRACE
Linear evaluation on learned node embeddings Linear evaluation on learned node embeddings
''' """
import numpy as np
import functools import functools
from sklearn.metrics import f1_score import numpy as np
from sklearn.linear_model import LogisticRegression from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split, GridSearchCV from sklearn.metrics import f1_score
from sklearn.model_selection import GridSearchCV, train_test_split
from sklearn.multiclass import OneVsRestClassifier from sklearn.multiclass import OneVsRestClassifier
from sklearn.preprocessing import normalize, OneHotEncoder from sklearn.preprocessing import OneHotEncoder, normalize
def repeat(n_times): def repeat(n_times):
...@@ -22,8 +22,9 @@ def repeat(n_times): ...@@ -22,8 +22,9 @@ def repeat(n_times):
for key in results[0].keys(): for key in results[0].keys():
values = [r[key] for r in results] values = [r[key] for r in results]
statistics[key] = { statistics[key] = {
'mean': np.mean(values), "mean": np.mean(values),
'std': np.std(values)} "std": np.std(values),
}
print_statistics(statistics, f.__name__) print_statistics(statistics, f.__name__)
return statistics return statistics
...@@ -41,41 +42,49 @@ def prob_to_one_hot(y_pred): ...@@ -41,41 +42,49 @@ def prob_to_one_hot(y_pred):
def print_statistics(statistics, function_name): def print_statistics(statistics, function_name):
print(f'(E) | {function_name}:', end=' ') print(f"(E) | {function_name}:", end=" ")
for i, key in enumerate(statistics.keys()): for i, key in enumerate(statistics.keys()):
mean = statistics[key]['mean'] mean = statistics[key]["mean"]
std = statistics[key]['std'] std = statistics[key]["std"]
print(f'{key}={mean:.4f}+-{std:.4f}', end='') print(f"{key}={mean:.4f}+-{std:.4f}", end="")
if i != len(statistics.keys()) - 1: if i != len(statistics.keys()) - 1:
print(',', end=' ') print(",", end=" ")
else: else:
print() print()
@repeat(3) @repeat(3)
def label_classification(embeddings, y, train_mask, test_mask, split='random', ratio=0.1): def label_classification(
embeddings, y, train_mask, test_mask, split="random", ratio=0.1
):
X = embeddings.detach().cpu().numpy() X = embeddings.detach().cpu().numpy()
Y = y.detach().cpu().numpy() Y = y.detach().cpu().numpy()
Y = Y.reshape(-1, 1) Y = Y.reshape(-1, 1)
onehot_encoder = OneHotEncoder(categories='auto').fit(Y) onehot_encoder = OneHotEncoder(categories="auto").fit(Y)
Y = onehot_encoder.transform(Y).toarray().astype(np.bool) Y = onehot_encoder.transform(Y).toarray().astype(np.bool)
X = normalize(X, norm='l2') X = normalize(X, norm="l2")
if split == 'random': if split == "random":
X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=1 - ratio) X_train, X_test, y_train, y_test = train_test_split(
elif split == 'public': X, Y, test_size=1 - ratio
)
elif split == "public":
X_train = X[train_mask] X_train = X[train_mask]
X_test = X[test_mask] X_test = X[test_mask]
y_train = Y[train_mask] y_train = Y[train_mask]
y_test = Y[test_mask] y_test = Y[test_mask]
logreg = LogisticRegression(solver='liblinear') logreg = LogisticRegression(solver="liblinear")
c = 2.0 ** np.arange(-10, 10) c = 2.0 ** np.arange(-10, 10)
clf = GridSearchCV(estimator=OneVsRestClassifier(logreg), clf = GridSearchCV(
param_grid=dict(estimator__C=c), n_jobs=8, cv=5, estimator=OneVsRestClassifier(logreg),
verbose=0) param_grid=dict(estimator__C=c),
n_jobs=8,
cv=5,
verbose=0,
)
clf.fit(X_train, y_train) clf.fit(X_train, y_train)
y_pred = clf.predict_proba(X_test) y_pred = clf.predict_proba(X_test)
...@@ -84,7 +93,4 @@ def label_classification(embeddings, y, train_mask, test_mask, split='random', r ...@@ -84,7 +93,4 @@ def label_classification(embeddings, y, train_mask, test_mask, split='random', r
micro = f1_score(y_test, y_pred, average="micro") micro = f1_score(y_test, y_pred, average="micro")
macro = f1_score(y_test, y_pred, average="macro") macro = f1_score(y_test, y_pred, average="macro")
return { return {"F1Mi": micro, "F1Ma": macro}
'F1Mi': micro,
'F1Ma': macro
}
import argparse import argparse
from model import Grace import warnings
from aug import aug
from dataset import load
import numpy as np import numpy as np
import torch as th import torch as th
import torch.nn as nn import torch.nn as nn
from aug import aug
from dataset import load
from eval import label_classification from eval import label_classification
import warnings from model import Grace
warnings.filterwarnings('ignore') warnings.filterwarnings("ignore")
def count_parameters(model): def count_parameters(model):
return sum([np.prod(p.size()) for p in model.parameters() if p.requires_grad]) return sum(
[np.prod(p.size()) for p in model.parameters() if p.requires_grad]
)
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--dataname', type=str, default='cora') parser.add_argument("--dataname", type=str, default="cora")
parser.add_argument('--gpu', type=int, default=0) parser.add_argument("--gpu", type=int, default=0)
parser.add_argument('--split', type=str, default='random') parser.add_argument("--split", type=str, default="random")
parser.add_argument('--epochs', type=int, default=500, help='Number of training periods.') parser.add_argument(
parser.add_argument('--lr', type=float, default=0.001, help='Learning rate.') "--epochs", type=int, default=500, help="Number of training periods."
parser.add_argument('--wd', type=float, default=1e-5, help='Weight decay.') )
parser.add_argument('--temp', type=float, default=1.0, help='Temperature.') parser.add_argument("--lr", type=float, default=0.001, help="Learning rate.")
parser.add_argument("--wd", type=float, default=1e-5, help="Weight decay.")
parser.add_argument('--act_fn', type=str, default='relu') parser.add_argument("--temp", type=float, default=1.0, help="Temperature.")
parser.add_argument("--hid_dim", type=int, default=256, help='Hidden layer dim.') parser.add_argument("--act_fn", type=str, default="relu")
parser.add_argument("--out_dim", type=int, default=256, help='Output layer dim.')
parser.add_argument(
parser.add_argument("--num_layers", type=int, default=2, help='Number of GNN layers.') "--hid_dim", type=int, default=256, help="Hidden layer dim."
parser.add_argument('--der1', type=float, default=0.2, help='Drop edge ratio of the 1st augmentation.') )
parser.add_argument('--der2', type=float, default=0.2, help='Drop edge ratio of the 2nd augmentation.') parser.add_argument(
parser.add_argument('--dfr1', type=float, default=0.2, help='Drop feature ratio of the 1st augmentation.') "--out_dim", type=int, default=256, help="Output layer dim."
parser.add_argument('--dfr2', type=float, default=0.2, help='Drop feature ratio of the 2nd augmentation.') )
parser.add_argument(
"--num_layers", type=int, default=2, help="Number of GNN layers."
)
parser.add_argument(
"--der1",
type=float,
default=0.2,
help="Drop edge ratio of the 1st augmentation.",
)
parser.add_argument(
"--der2",
type=float,
default=0.2,
help="Drop edge ratio of the 2nd augmentation.",
)
parser.add_argument(
"--dfr1",
type=float,
default=0.2,
help="Drop feature ratio of the 1st augmentation.",
)
parser.add_argument(
"--dfr2",
type=float,
default=0.2,
help="Drop feature ratio of the 2nd augmentation.",
)
args = parser.parse_args() args = parser.parse_args()
if args.gpu != -1 and th.cuda.is_available(): if args.gpu != -1 and th.cuda.is_available():
args.device = 'cuda:{}'.format(args.gpu) args.device = "cuda:{}".format(args.gpu)
else: else:
args.device = 'cpu' args.device = "cpu"
if __name__ == '__main__': if __name__ == "__main__":
# Step 1: Load hyperparameters =================================================================== # # Step 1: Load hyperparameters =================================================================== #
lr = args.lr lr = args.lr
...@@ -53,7 +82,7 @@ if __name__ == '__main__': ...@@ -53,7 +82,7 @@ if __name__ == '__main__':
out_dim = args.out_dim out_dim = args.out_dim
num_layers = args.num_layers num_layers = args.num_layers
act_fn = ({'relu': nn.ReLU(), 'prelu': nn.PReLU()})[args.act_fn] act_fn = ({"relu": nn.ReLU(), "prelu": nn.PReLU()})[args.act_fn]
drop_edge_rate_1 = args.der1 drop_edge_rate_1 = args.der1
drop_edge_rate_2 = args.der2 drop_edge_rate_2 = args.der2
...@@ -71,7 +100,7 @@ if __name__ == '__main__': ...@@ -71,7 +100,7 @@ if __name__ == '__main__':
# Step 3: Create model =================================================================== # # Step 3: Create model =================================================================== #
model = Grace(in_dim, hid_dim, out_dim, num_layers, act_fn, temp) model = Grace(in_dim, hid_dim, out_dim, num_layers, act_fn, temp)
model = model.to(args.device) model = model.to(args.device)
print(f'# params: {count_parameters(model)}') print(f"# params: {count_parameters(model)}")
optimizer = th.optim.Adam(model.parameters(), lr=lr, weight_decay=wd) optimizer = th.optim.Adam(model.parameters(), lr=lr, weight_decay=wd)
...@@ -92,7 +121,7 @@ if __name__ == '__main__': ...@@ -92,7 +121,7 @@ if __name__ == '__main__':
loss.backward() loss.backward()
optimizer.step() optimizer.step()
print(f'Epoch={epoch:03d}, loss={loss.item():.4f}') print(f"Epoch={epoch:03d}, loss={loss.item():.4f}")
# Step 5: Linear evaluation ============================================================== # # Step 5: Linear evaluation ============================================================== #
print("=== Final ===") print("=== Final ===")
...@@ -102,5 +131,7 @@ if __name__ == '__main__': ...@@ -102,5 +131,7 @@ if __name__ == '__main__':
feat = feat.to(args.device) feat = feat.to(args.device)
embeds = model.get_embedding(graph, feat) embeds = model.get_embedding(graph, feat)
'''Evaluation Embeddings ''' """Evaluation Embeddings """
label_classification(embeds, labels, train_mask, test_mask, split=args.split) label_classification(
\ No newline at end of file embeds, labels, train_mask, test_mask, split=args.split
)
import torch as th import torch as th
import torch.nn as nn import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
from dgl.nn import GraphConv from dgl.nn import GraphConv
# Multi-layer Graph Convolutional Networks # Multi-layer Graph Convolutional Networks
class GCN(nn.Module): class GCN(nn.Module):
def __init__(self, in_dim, out_dim, act_fn, num_layers = 2): def __init__(self, in_dim, out_dim, act_fn, num_layers=2):
super(GCN, self).__init__() super(GCN, self).__init__()
assert num_layers >= 2 assert num_layers >= 2
...@@ -26,6 +27,7 @@ class GCN(nn.Module): ...@@ -26,6 +27,7 @@ class GCN(nn.Module):
return feat return feat
# Multi-layer(2-layer) Perceptron # Multi-layer(2-layer) Perceptron
class MLP(nn.Module): class MLP(nn.Module):
def __init__(self, in_dim, out_dim): def __init__(self, in_dim, out_dim):
...@@ -56,6 +58,7 @@ class Grace(nn.Module): ...@@ -56,6 +58,7 @@ class Grace(nn.Module):
temp: float temp: float
Temperature constant. Temperature constant.
""" """
def __init__(self, in_dim, hid_dim, out_dim, num_layers, act_fn, temp): def __init__(self, in_dim, hid_dim, out_dim, num_layers, act_fn, temp):
super(Grace, self).__init__() super(Grace, self).__init__()
self.encoder = GCN(in_dim, hid_dim, act_fn, num_layers) self.encoder = GCN(in_dim, hid_dim, act_fn, num_layers)
...@@ -74,8 +77,8 @@ class Grace(nn.Module): ...@@ -74,8 +77,8 @@ class Grace(nn.Module):
# calculate SimCLR loss # calculate SimCLR loss
f = lambda x: th.exp(x / self.temp) f = lambda x: th.exp(x / self.temp)
refl_sim = f(self.sim(z1, z1)) # intra-view pairs refl_sim = f(self.sim(z1, z1)) # intra-view pairs
between_sim = f(self.sim(z1, z2)) # inter-view pairs between_sim = f(self.sim(z1, z2)) # inter-view pairs
# between_sim.diag(): positive pairs # between_sim.diag(): positive pairs
x1 = refl_sim.sum(1) + between_sim.sum(1) - refl_sim.diag() x1 = refl_sim.sum(1) + between_sim.sum(1) - refl_sim.diag()
...@@ -104,4 +107,4 @@ class Grace(nn.Module): ...@@ -104,4 +107,4 @@ class Grace(nn.Module):
ret = (l1 + l2) * 0.5 ret = (l1 + l2) * 0.5
return ret.mean() return ret.mean()
\ No newline at end of file
import argparse import argparse
import warnings
import numpy as np import numpy as np
import torch as th import torch as th
import torch.optim as optim
import torch.nn as nn import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
import torch.optim as optim
from model import GRAND
import dgl import dgl
from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
from model import GRAND warnings.filterwarnings("ignore")
import warnings
warnings.filterwarnings('ignore')
def argument(): def argument():
parser = argparse.ArgumentParser(description='GRAND') parser = argparse.ArgumentParser(description="GRAND")
# data source params # data source params
parser.add_argument('--dataname', type=str, default='cora', help='Name of dataset.') parser.add_argument(
"--dataname", type=str, default="cora", help="Name of dataset."
)
# cuda params # cuda params
parser.add_argument('--gpu', type=int, default=-1, help='GPU index. Default: -1, using CPU.') parser.add_argument(
"--gpu", type=int, default=-1, help="GPU index. Default: -1, using CPU."
)
# training params # training params
parser.add_argument('--epochs', type=int, default=200, help='Training epochs.') parser.add_argument(
parser.add_argument('--early_stopping', type=int, default=200, help='Patient epochs to wait before early stopping.') "--epochs", type=int, default=200, help="Training epochs."
parser.add_argument('--lr', type=float, default=0.01, help='Learning rate.') )
parser.add_argument('--weight_decay', type=float, default=5e-4, help='L2 reg.') parser.add_argument(
"--early_stopping",
type=int,
default=200,
help="Patient epochs to wait before early stopping.",
)
parser.add_argument("--lr", type=float, default=0.01, help="Learning rate.")
parser.add_argument(
"--weight_decay", type=float, default=5e-4, help="L2 reg."
)
# model params # model params
parser.add_argument("--hid_dim", type=int, default=32, help='Hidden layer dimensionalities.') parser.add_argument(
parser.add_argument('--dropnode_rate', type=float, default=0.5, "--hid_dim", type=int, default=32, help="Hidden layer dimensionalities."
help='Dropnode rate (1 - keep probability).') )
parser.add_argument('--input_droprate', type=float, default=0.0, parser.add_argument(
help='dropout rate of input layer') "--dropnode_rate",
parser.add_argument('--hidden_droprate', type=float, default=0.0, type=float,
help='dropout rate of hidden layer') default=0.5,
parser.add_argument('--order', type=int, default=8, help='Propagation step') help="Dropnode rate (1 - keep probability).",
parser.add_argument('--sample', type=int, default=4, help='Sampling times of dropnode') )
parser.add_argument('--tem', type=float, default=0.5, help='Sharpening temperature') parser.add_argument(
parser.add_argument('--lam', type=float, default=1., help='Coefficient of consistency regularization') "--input_droprate",
parser.add_argument('--use_bn', action='store_true', default=False, help='Using Batch Normalization') type=float,
default=0.0,
help="dropout rate of input layer",
)
parser.add_argument(
"--hidden_droprate",
type=float,
default=0.0,
help="dropout rate of hidden layer",
)
parser.add_argument("--order", type=int, default=8, help="Propagation step")
parser.add_argument(
"--sample", type=int, default=4, help="Sampling times of dropnode"
)
parser.add_argument(
"--tem", type=float, default=0.5, help="Sharpening temperature"
)
parser.add_argument(
"--lam",
type=float,
default=1.0,
help="Coefficient of consistency regularization",
)
parser.add_argument(
"--use_bn",
action="store_true",
default=False,
help="Using Batch Normalization",
)
args = parser.parse_args() args = parser.parse_args()
# check cuda # check cuda
if args.gpu != -1 and th.cuda.is_available(): if args.gpu != -1 and th.cuda.is_available():
args.device = 'cuda:{}'.format(args.gpu) args.device = "cuda:{}".format(args.gpu)
else: else:
args.device = 'cpu' args.device = "cpu"
return args return args
def consis_loss(logps, temp, lam): def consis_loss(logps, temp, lam):
ps = [th.exp(p) for p in logps] ps = [th.exp(p) for p in logps]
ps = th.stack(ps, dim = 2) ps = th.stack(ps, dim=2)
avg_p = th.mean(ps, dim = 2) avg_p = th.mean(ps, dim=2)
sharp_p = (th.pow(avg_p, 1./temp) / th.sum(th.pow(avg_p, 1./temp), dim=1, keepdim=True)).detach() sharp_p = (
th.pow(avg_p, 1.0 / temp)
/ th.sum(th.pow(avg_p, 1.0 / temp), dim=1, keepdim=True)
).detach()
sharp_p = sharp_p.unsqueeze(2) sharp_p = sharp_p.unsqueeze(2)
loss = th.mean(th.sum(th.pow(ps - sharp_p, 2), dim = 1, keepdim=True)) loss = th.mean(th.sum(th.pow(ps - sharp_p, 2), dim=1, keepdim=True))
loss = lam * loss loss = lam * loss
return loss return loss
if __name__ == '__main__':
if __name__ == "__main__":
# Step 1: Prepare graph data and retrieve train/validation/test index ============================= # # Step 1: Prepare graph data and retrieve train/validation/test index ============================= #
# Load from DGL dataset # Load from DGL dataset
args = argument() args = argument()
print(args) print(args)
if args.dataname == 'cora': if args.dataname == "cora":
dataset = CoraGraphDataset() dataset = CoraGraphDataset()
elif args.dataname == 'citeseer': elif args.dataname == "citeseer":
dataset = CiteseerGraphDataset() dataset = CiteseerGraphDataset()
elif args.dataname == 'pubmed': elif args.dataname == "pubmed":
dataset = PubmedGraphDataset() dataset = PubmedGraphDataset()
graph = dataset[0] graph = dataset[0]
graph = dgl.add_self_loop(graph) graph = dgl.add_self_loop(graph)
device = args.device device = args.device
...@@ -86,100 +133,121 @@ if __name__ == '__main__': ...@@ -86,100 +133,121 @@ if __name__ == '__main__':
n_classes = dataset.num_classes n_classes = dataset.num_classes
# retrieve labels of ground truth # retrieve labels of ground truth
labels = graph.ndata.pop('label').to(device).long() labels = graph.ndata.pop("label").to(device).long()
# Extract node features # Extract node features
feats = graph.ndata.pop('feat').to(device) feats = graph.ndata.pop("feat").to(device)
n_features = feats.shape[-1] n_features = feats.shape[-1]
# retrieve masks for train/validation/test # retrieve masks for train/validation/test
train_mask = graph.ndata.pop('train_mask') train_mask = graph.ndata.pop("train_mask")
val_mask = graph.ndata.pop('val_mask') val_mask = graph.ndata.pop("val_mask")
test_mask = graph.ndata.pop('test_mask') test_mask = graph.ndata.pop("test_mask")
train_idx = th.nonzero(train_mask, as_tuple=False).squeeze().to(device) train_idx = th.nonzero(train_mask, as_tuple=False).squeeze().to(device)
val_idx = th.nonzero(val_mask, as_tuple=False).squeeze().to(device) val_idx = th.nonzero(val_mask, as_tuple=False).squeeze().to(device)
test_idx = th.nonzero(test_mask, as_tuple=False).squeeze().to(device) test_idx = th.nonzero(test_mask, as_tuple=False).squeeze().to(device)
# Step 2: Create model =================================================================== # # Step 2: Create model =================================================================== #
model = GRAND(n_features, args.hid_dim, n_classes, args.sample, args.order, model = GRAND(
args.dropnode_rate, args.input_droprate, n_features,
args.hidden_droprate, args.use_bn) args.hid_dim,
n_classes,
args.sample,
args.order,
args.dropnode_rate,
args.input_droprate,
args.hidden_droprate,
args.use_bn,
)
model = model.to(args.device) model = model.to(args.device)
graph = graph.to(args.device) graph = graph.to(args.device)
# Step 3: Create training components ===================================================== # # Step 3: Create training components ===================================================== #
loss_fn = nn.NLLLoss() loss_fn = nn.NLLLoss()
opt = optim.Adam(model.parameters(), lr = args.lr, weight_decay = args.weight_decay) opt = optim.Adam(
model.parameters(), lr=args.lr, weight_decay=args.weight_decay
)
loss_best = np.inf loss_best = np.inf
acc_best = 0 acc_best = 0
# Step 4: training epoches =============================================================== # # Step 4: training epoches =============================================================== #
for epoch in range(args.epochs): for epoch in range(args.epochs):
''' Training ''' """Training"""
model.train() model.train()
loss_sup = 0 loss_sup = 0
logits = model(graph, feats, True) logits = model(graph, feats, True)
# calculate supervised loss # calculate supervised loss
for k in range(args.sample): for k in range(args.sample):
loss_sup += F.nll_loss(logits[k][train_idx], labels[train_idx]) loss_sup += F.nll_loss(logits[k][train_idx], labels[train_idx])
loss_sup = loss_sup/args.sample loss_sup = loss_sup / args.sample
# calculate consistency loss # calculate consistency loss
loss_consis = consis_loss(logits, args.tem, args.lam) loss_consis = consis_loss(logits, args.tem, args.lam)
loss_train = loss_sup + loss_consis loss_train = loss_sup + loss_consis
acc_train = th.sum(logits[0][train_idx].argmax(dim=1) == labels[train_idx]).item() / len(train_idx) acc_train = th.sum(
logits[0][train_idx].argmax(dim=1) == labels[train_idx]
).item() / len(train_idx)
# backward # backward
opt.zero_grad() opt.zero_grad()
loss_train.backward() loss_train.backward()
opt.step() opt.step()
''' Validating ''' """ Validating """
model.eval() model.eval()
with th.no_grad(): with th.no_grad():
val_logits = model(graph, feats, False) val_logits = model(graph, feats, False)
loss_val = F.nll_loss(val_logits[val_idx], labels[val_idx]) loss_val = F.nll_loss(val_logits[val_idx], labels[val_idx])
acc_val = th.sum(val_logits[val_idx].argmax(dim=1) == labels[val_idx]).item() / len(val_idx) acc_val = th.sum(
val_logits[val_idx].argmax(dim=1) == labels[val_idx]
).item() / len(val_idx)
# Print out performance # Print out performance
print("In epoch {}, Train Acc: {:.4f} | Train Loss: {:.4f} ,Val Acc: {:.4f} | Val Loss: {:.4f}". print(
format(epoch, acc_train, loss_train.item(), acc_val, loss_val.item())) "In epoch {}, Train Acc: {:.4f} | Train Loss: {:.4f} ,Val Acc: {:.4f} | Val Loss: {:.4f}".format(
epoch,
acc_train,
loss_train.item(),
acc_val,
loss_val.item(),
)
)
# set early stopping counter # set early stopping counter
if loss_val < loss_best or acc_val > acc_best: if loss_val < loss_best or acc_val > acc_best:
if loss_val < loss_best: if loss_val < loss_best:
best_epoch = epoch best_epoch = epoch
th.save(model.state_dict(), args.dataname +'.pkl') th.save(model.state_dict(), args.dataname + ".pkl")
no_improvement = 0 no_improvement = 0
loss_best = min(loss_val, loss_best) loss_best = min(loss_val, loss_best)
acc_best = max(acc_val, acc_best) acc_best = max(acc_val, acc_best)
else: else:
no_improvement += 1 no_improvement += 1
if no_improvement == args.early_stopping: if no_improvement == args.early_stopping:
print('Early stopping.') print("Early stopping.")
break break
print("Optimization Finished!") print("Optimization Finished!")
print('Loading {}th epoch'.format(best_epoch)) print("Loading {}th epoch".format(best_epoch))
model.load_state_dict(th.load(args.dataname +'.pkl')) model.load_state_dict(th.load(args.dataname + ".pkl"))
''' Testing ''' """ Testing """
model.eval() model.eval()
test_logits = model(graph, feats, False)
test_acc = th.sum(test_logits[test_idx].argmax(dim=1) == labels[test_idx]).item() / len(test_idx)
print("Test Acc: {:.4f}".format(test_acc)) test_logits = model(graph, feats, False)
test_acc = th.sum(
test_logits[test_idx].argmax(dim=1) == labels[test_idx]
).item() / len(test_idx)
print("Test Acc: {:.4f}".format(test_acc))
import numpy as np import numpy as np
import torch as th import torch as th
import torch.nn as nn import torch.nn as nn
import dgl.function as fn
import torch.nn.functional as F import torch.nn.functional as F
import dgl.function as fn
def drop_node(feats, drop_rate, training): def drop_node(feats, drop_rate, training):
n = feats.shape[0] n = feats.shape[0]
drop_rates = th.FloatTensor(np.ones(n) * drop_rate) drop_rates = th.FloatTensor(np.ones(n) * drop_rate)
if training: if training:
masks = th.bernoulli(1. - drop_rates).unsqueeze(1) masks = th.bernoulli(1.0 - drop_rates).unsqueeze(1)
feats = masks.to(feats.device) * feats feats = masks.to(feats.device) * feats
else: else:
feats = feats * (1. - drop_rate) feats = feats * (1.0 - drop_rate)
return feats return feats
class MLP(nn.Module): class MLP(nn.Module):
def __init__(self, nfeat, nhid, nclass, input_droprate, hidden_droprate, use_bn =False): def __init__(
self, nfeat, nhid, nclass, input_droprate, hidden_droprate, use_bn=False
):
super(MLP, self).__init__() super(MLP, self).__init__()
self.layer1 = nn.Linear(nfeat, nhid, bias = True) self.layer1 = nn.Linear(nfeat, nhid, bias=True)
self.layer2 = nn.Linear(nhid, nclass, bias = True) self.layer2 = nn.Linear(nhid, nclass, bias=True)
self.input_dropout = nn.Dropout(input_droprate) self.input_dropout = nn.Dropout(input_droprate)
self.hidden_dropout = nn.Dropout(hidden_droprate) self.hidden_dropout = nn.Dropout(hidden_droprate)
self.bn1 = nn.BatchNorm1d(nfeat) self.bn1 = nn.BatchNorm1d(nfeat)
self.bn2 = nn.BatchNorm1d(nhid) self.bn2 = nn.BatchNorm1d(nhid)
self.use_bn = use_bn self.use_bn = use_bn
def reset_parameters(self): def reset_parameters(self):
self.layer1.reset_parameters() self.layer1.reset_parameters()
self.layer2.reset_parameters() self.layer2.reset_parameters()
def forward(self, x): def forward(self, x):
if self.use_bn: if self.use_bn:
x = self.bn1(x) x = self.bn1(x)
x = self.input_dropout(x) x = self.input_dropout(x)
x = F.relu(self.layer1(x)) x = F.relu(self.layer1(x))
if self.use_bn: if self.use_bn:
x = self.bn2(x) x = self.bn2(x)
x = self.hidden_dropout(x) x = self.hidden_dropout(x)
x = self.layer2(x) x = self.layer2(x)
return x return x
def GRANDConv(graph, feats, order): def GRANDConv(graph, feats, order):
''' """
Parameters Parameters
----------- -----------
graph: dgl.Graph graph: dgl.Graph
The input graph The input graph
feats: Tensor (n_nodes * feat_dim) feats: Tensor (n_nodes * feat_dim)
Node features Node features
order: int order: int
Propagation Steps Propagation Steps
''' """
with graph.local_scope(): with graph.local_scope():
''' Calculate Symmetric normalized adjacency matrix \hat{A} ''' """Calculate Symmetric normalized adjacency matrix \hat{A}"""
degs = graph.in_degrees().float().clamp(min=1) degs = graph.in_degrees().float().clamp(min=1)
norm = th.pow(degs, -0.5).to(feats.device).unsqueeze(1) norm = th.pow(degs, -0.5).to(feats.device).unsqueeze(1)
graph.ndata['norm'] = norm graph.ndata["norm"] = norm
graph.apply_edges(fn.u_mul_v('norm', 'norm', 'weight')) graph.apply_edges(fn.u_mul_v("norm", "norm", "weight"))
''' Graph Conv ''' """ Graph Conv """
x = feats x = feats
y = 0+feats y = 0 + feats
for i in range(order): for i in range(order):
graph.ndata['h'] = x graph.ndata["h"] = x
graph.update_all(fn.u_mul_e('h', 'weight', 'm'), fn.sum('m', 'h')) graph.update_all(fn.u_mul_e("h", "weight", "m"), fn.sum("m", "h"))
x = graph.ndata.pop('h') x = graph.ndata.pop("h")
y.add_(x) y.add_(x)
return y /(order + 1) return y / (order + 1)
class GRAND(nn.Module): class GRAND(nn.Module):
r""" r"""
...@@ -108,16 +114,19 @@ class GRAND(nn.Module): ...@@ -108,16 +114,19 @@ class GRAND(nn.Module):
If True, use batch normalization. If True, use batch normalization.
""" """
def __init__(self,
in_dim, def __init__(
hid_dim, self,
n_class, in_dim,
S = 1, hid_dim,
K = 3, n_class,
node_dropout=0.0, S=1,
input_droprate = 0.0, K=3,
hidden_droprate = 0.0, node_dropout=0.0,
batchnorm=False): input_droprate=0.0,
hidden_droprate=0.0,
batchnorm=False,
):
super(GRAND, self).__init__() super(GRAND, self).__init__()
self.in_dim = in_dim self.in_dim = in_dim
...@@ -125,29 +134,31 @@ class GRAND(nn.Module): ...@@ -125,29 +134,31 @@ class GRAND(nn.Module):
self.S = S self.S = S
self.K = K self.K = K
self.n_class = n_class self.n_class = n_class
self.mlp = MLP(in_dim, hid_dim, n_class, input_droprate, hidden_droprate, batchnorm) self.mlp = MLP(
in_dim, hid_dim, n_class, input_droprate, hidden_droprate, batchnorm
)
self.dropout = node_dropout self.dropout = node_dropout
self.node_dropout = nn.Dropout(node_dropout) self.node_dropout = nn.Dropout(node_dropout)
def forward(self, graph, feats, training = True): def forward(self, graph, feats, training=True):
X = feats X = feats
S = self.S S = self.S
if training: # Training Mode if training: # Training Mode
output_list = [] output_list = []
for s in range(S): for s in range(S):
drop_feat = drop_node(X, self.dropout, True) # Drop node drop_feat = drop_node(X, self.dropout, True) # Drop node
feat = GRANDConv(graph, drop_feat, self.K) # Graph Convolution feat = GRANDConv(graph, drop_feat, self.K) # Graph Convolution
output_list.append(th.log_softmax(self.mlp(feat), dim=-1)) # Prediction output_list.append(
th.log_softmax(self.mlp(feat), dim=-1)
) # Prediction
return output_list return output_list
else: # Inference Mode else: # Inference Mode
drop_feat = drop_node(X, self.dropout, False) drop_feat = drop_node(X, self.dropout, False)
X = GRANDConv(graph, drop_feat, self.K) X = GRANDConv(graph, drop_feat, self.K)
return th.log_softmax(self.mlp(X), dim = -1) return th.log_softmax(self.mlp(X), dim=-1)
from ged import graph_edit_distance
import dgl
import numpy as np import numpy as np
from ged import graph_edit_distance
import dgl
src1 = [0, 1, 2, 3, 4, 5]; src1 = [0, 1, 2, 3, 4, 5]
dst1 = [1, 2, 3, 4, 5, 6]; dst1 = [1, 2, 3, 4, 5, 6]
src2 = [0, 1, 3, 4, 5]; src2 = [0, 1, 3, 4, 5]
dst2 = [1, 2, 4, 5, 6]; dst2 = [1, 2, 4, 5, 6]
G1 = dgl.DGLGraph((src1, dst1)) G1 = dgl.DGLGraph((src1, dst1))
...@@ -15,55 +15,73 @@ G2 = dgl.DGLGraph((src2, dst2)) ...@@ -15,55 +15,73 @@ G2 = dgl.DGLGraph((src2, dst2))
# Exact edit distance with astar search # Exact edit distance with astar search
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G1, algorithm='astar') distance, node_mapping, edge_mapping = graph_edit_distance(
print(distance) # 0.0 G1, G1, algorithm="astar"
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='astar') )
print(distance) # 1.0 print(distance) # 0.0
distance, node_mapping, edge_mapping = graph_edit_distance(
G1, G2, algorithm="astar"
)
print(distance) # 1.0
# With user-input cost matrices # With user-input cost matrices
node_substitution_cost = np.empty((G1.number_of_nodes(), G2.number_of_nodes())); node_substitution_cost = np.empty((G1.number_of_nodes(), G2.number_of_nodes()))
G1_node_deletion_cost = np.empty(G1.number_of_nodes()); G1_node_deletion_cost = np.empty(G1.number_of_nodes())
G2_node_insertion_cost = np.empty(G2.number_of_nodes()); G2_node_insertion_cost = np.empty(G2.number_of_nodes())
edge_substitution_cost = np.empty((G1.number_of_edges(), G2.number_of_edges())); edge_substitution_cost = np.empty((G1.number_of_edges(), G2.number_of_edges()))
G1_edge_deletion_cost = np.empty(G1.number_of_edges()); G1_edge_deletion_cost = np.empty(G1.number_of_edges())
G2_edge_insertion_cost = np.empty(G2.number_of_edges()); G2_edge_insertion_cost = np.empty(G2.number_of_edges())
# Node substitution cost of 0 when node-ids are same, else 1 # Node substitution cost of 0 when node-ids are same, else 1
node_substitution_cost.fill(1.0); node_substitution_cost.fill(1.0)
for i in range(G1.number_of_nodes()): for i in range(G1.number_of_nodes()):
for j in range(G2.number_of_nodes()): for j in range(G2.number_of_nodes()):
node_substitution_cost[i,j] = 0.0; node_substitution_cost[i, j] = 0.0
# Node insertion/deletion cost of 1 # Node insertion/deletion cost of 1
G1_node_deletion_cost.fill(1.0); G1_node_deletion_cost.fill(1.0)
G2_node_insertion_cost.fill(1.0); G2_node_insertion_cost.fill(1.0)
# Edge substitution cost of 0 # Edge substitution cost of 0
edge_substitution_cost.fill(0.0); edge_substitution_cost.fill(0.0)
# Edge insertion/deletion cost of 0.5 # Edge insertion/deletion cost of 0.5
G1_edge_deletion_cost.fill(0.5); G1_edge_deletion_cost.fill(0.5)
G2_edge_insertion_cost.fill(0.5); G2_edge_insertion_cost.fill(0.5)
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, \ distance, node_mapping, edge_mapping = graph_edit_distance(
node_substitution_cost, edge_substitution_cost, \ G1,
G1_node_deletion_cost, G2_node_insertion_cost, \ G2,
G1_edge_deletion_cost, G2_edge_insertion_cost, \ node_substitution_cost,
algorithm="astar") edge_substitution_cost,
G1_node_deletion_cost,
G2_node_insertion_cost,
G1_edge_deletion_cost,
G2_edge_insertion_cost,
algorithm="astar",
)
print(distance) #0.5 print(distance) # 0.5
# Approximate edit distance with beam search, it is more than or equal to the exact edit distance # Approximate edit distance with beam search, it is more than or equal to the exact edit distance
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='beam', max_beam_size=2) distance, node_mapping, edge_mapping = graph_edit_distance(
print(distance) # 3.0 G1, G2, algorithm="beam", max_beam_size=2
)
print(distance) # 3.0
# Approximate edit distance with bipartite heuristic, it is more than or equal to the exact edit distance # Approximate edit distance with bipartite heuristic, it is more than or equal to the exact edit distance
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='bipartite') distance, node_mapping, edge_mapping = graph_edit_distance(
print(distance) # 9.0, can be different as multiple solutions possible for the intermediate LAP used in this approximation G1, G2, algorithm="bipartite"
)
print(
distance
) # 9.0, can be different as multiple solutions possible for the intermediate LAP used in this approximation
# Approximate edit distance with hausdorff heuristic, it is less than or equal to the exact edit distance # Approximate edit distance with hausdorff heuristic, it is less than or equal to the exact edit distance
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='hausdorff') distance, node_mapping, edge_mapping = graph_edit_distance(
print(distance) # 0.0 G1, G2, algorithm="hausdorff"
\ No newline at end of file )
print(distance) # 0.0
import sklearn.linear_model as lm
import sklearn.metrics as skm
import torch as th import torch as th
import torch.nn as nn
import torch.functional as F import torch.functional as F
import torch.nn as nn
import tqdm
import dgl import dgl
import dgl.nn as dglnn import dgl.nn as dglnn
import sklearn.linear_model as lm
import sklearn.metrics as skm
import tqdm
class SAGE(nn.Module): class SAGE(nn.Module):
def __init__(self, in_feats, n_hidden, n_classes, n_layers, activation, dropout): def __init__(
self, in_feats, n_hidden, n_classes, n_layers, activation, dropout
):
super().__init__() super().__init__()
self.init(in_feats, n_hidden, n_classes, n_layers, activation, dropout) self.init(in_feats, n_hidden, n_classes, n_layers, activation, dropout)
def init(self, in_feats, n_hidden, n_classes, n_layers, activation, dropout): def init(
self, in_feats, n_hidden, n_classes, n_layers, activation, dropout
):
self.n_layers = n_layers self.n_layers = n_layers
self.n_hidden = n_hidden self.n_hidden = n_hidden
self.n_classes = n_classes self.n_classes = n_classes
self.layers = nn.ModuleList() self.layers = nn.ModuleList()
if n_layers > 1: if n_layers > 1:
self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, 'mean')) self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, "mean"))
for i in range(1, n_layers - 1): for i in range(1, n_layers - 1):
self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, 'mean')) self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, "mean"))
self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, 'mean')) self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, "mean"))
else: else:
self.layers.append(dglnn.SAGEConv(in_feats, n_classes, 'mean')) self.layers.append(dglnn.SAGEConv(in_feats, n_classes, "mean"))
self.dropout = nn.Dropout(dropout) self.dropout = nn.Dropout(dropout)
self.activation = activation self.activation = activation
...@@ -51,7 +57,10 @@ class SAGE(nn.Module): ...@@ -51,7 +57,10 @@ class SAGE(nn.Module):
# on each layer are of course splitted in batches. # on each layer are of course splitted in batches.
# TODO: can we standardize this? # TODO: can we standardize this?
for l, layer in enumerate(self.layers): for l, layer in enumerate(self.layers):
y = th.zeros(g.num_nodes(), self.n_hidden if l != len(self.layers) - 1 else self.n_classes) y = th.zeros(
g.num_nodes(),
self.n_hidden if l != len(self.layers) - 1 else self.n_classes,
)
sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1) sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
dataloader = dgl.dataloading.DataLoader( dataloader = dgl.dataloading.DataLoader(
...@@ -62,7 +71,8 @@ class SAGE(nn.Module): ...@@ -62,7 +71,8 @@ class SAGE(nn.Module):
batch_size=batch_size, batch_size=batch_size,
shuffle=False, shuffle=False,
drop_last=False, drop_last=False,
num_workers=num_workers) num_workers=num_workers,
)
for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader): for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
block = blocks[0] block = blocks[0]
...@@ -79,6 +89,7 @@ class SAGE(nn.Module): ...@@ -79,6 +89,7 @@ class SAGE(nn.Module):
x = y x = y
return y return y
def compute_acc_unsupervised(emb, labels, train_nids, val_nids, test_nids): def compute_acc_unsupervised(emb, labels, train_nids, val_nids, test_nids):
""" """
Compute the accuracy of prediction given the labels. Compute the accuracy of prediction given the labels.
...@@ -94,10 +105,10 @@ def compute_acc_unsupervised(emb, labels, train_nids, val_nids, test_nids): ...@@ -94,10 +105,10 @@ def compute_acc_unsupervised(emb, labels, train_nids, val_nids, test_nids):
emb = (emb - emb.mean(0, keepdims=True)) / emb.std(0, keepdims=True) emb = (emb - emb.mean(0, keepdims=True)) / emb.std(0, keepdims=True)
lr = lm.LogisticRegression(multi_class='multinomial', max_iter=10000) lr = lm.LogisticRegression(multi_class="multinomial", max_iter=10000)
lr.fit(emb[train_nids], train_labels) lr.fit(emb[train_nids], train_labels)
pred = lr.predict(emb) pred = lr.predict(emb)
f1_micro_eval = skm.f1_score(val_labels, pred[val_nids], average='micro') f1_micro_eval = skm.f1_score(val_labels, pred[val_nids], average="micro")
f1_micro_test = skm.f1_score(test_labels, pred[test_nids], average='micro') f1_micro_test = skm.f1_score(test_labels, pred[test_nids], average="micro")
return f1_micro_eval, f1_micro_test return f1_micro_eval, f1_micro_test
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment