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
from dgl.data import PPIDataset
import collections
#implement the collate_fn for dgl graph data class
PPIBatch = collections.namedtuple('PPIBatch', ['graph', 'label'])
# implement the collate_fn for dgl graph data class
PPIBatch = collections.namedtuple("PPIBatch", ["graph", "label"])
def batcher(device):
def batcher_dev(batch):
batch_graphs = dgl.batch(batch)
return PPIBatch(graph=batch_graphs,
label=batch_graphs.ndata['label'].to(device))
return PPIBatch(
graph=batch_graphs, label=batch_graphs.ndata["label"].to(device)
)
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'):
train_set = PPIDataset(mode='train')
valid_set = PPIDataset(mode='valid')
test_set = PPIDataset(mode='test')
#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
# 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"):
train_set = PPIDataset(mode="train")
valid_set = PPIDataset(mode="valid")
test_set = PPIDataset(mode="test")
# 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)):
g = dgl.heterograph({
('_N','_E','_N'): train_set[i].edges(),
('_N', 'self', '_N'): (train_set[i].nodes(), train_set[i].nodes())
})
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']
g = dgl.heterograph(
{
("_N", "_E", "_N"): train_set[i].edges(),
("_N", "self", "_N"): (
train_set[i].nodes(),
train_set[i].nodes(),
),
}
)
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
for i in range(len(valid_set)):
g = dgl.heterograph({
('_N','_E','_N'): valid_set[i].edges(),
('_N', 'self', '_N'): (valid_set[i].nodes(), valid_set[i].nodes())
})
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
g = dgl.heterograph(
{
("_N", "_E", "_N"): valid_set[i].edges(),
("_N", "self", "_N"): (
valid_set[i].nodes(),
valid_set[i].nodes(),
),
}
)
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)):
g = dgl.heterograph({
('_N','_E','_N'): test_set[i].edges(),
('_N', 'self', '_N'): (test_set[i].nodes(), test_set[i].nodes())
})
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
g = dgl.heterograph(
{
("_N", "_E", "_N"): test_set[i].edges(),
("_N", "self", "_N"): (
test_set[i].nodes(),
test_set[i].nodes(),
),
}
)
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
in_size = train_set[0].ndata['feat'].shape[1]
out_size = train_set[0].ndata['label'].shape[1]
in_size = train_set[0].ndata["feat"].shape[1]
out_size = train_set[0].ndata["label"].shape[1]
#prepare train, valid, and test dataloaders
train_loader = DataLoader(train_set, 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)
# prepare train, valid, and test dataloaders
train_loader = DataLoader(
train_set,
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
import argparse
import os
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from data_loader import load_PPI
from utils import evaluate_f1_score
import dgl
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):
def __init__(self, in_size, out_size, etypes, dropout=0.1):
......@@ -15,78 +18,112 @@ class GNNFiLMLayer(nn.Module):
self.in_size = in_size
self.out_size = out_size
#weights for different types of edges
self.W = nn.ModuleDict({
name : nn.Linear(in_size, out_size, bias = False) for name in etypes
})
# weights for different types of edges
self.W = nn.ModuleDict(
{name: nn.Linear(in_size, out_size, bias=False) for name in etypes}
)
#hypernets to learn the affine functions for different types of edges
self.film = nn.ModuleDict({
name : nn.Linear(in_size, 2*out_size, bias = False) for name in etypes
})
# hypernets to learn the affine functions for different types of edges
self.film = nn.ModuleDict(
{
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)
#dropout layer
# dropout layer
self.dropout = nn.Dropout(dropout)
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
#for each type of edges, compute messages and reduce them all
funcs = {} # message and reduce functions dict
# for each type of edges, compute messages and reduce them all
for srctype, etype, dsttype in g.canonical_etypes:
messages = self.W[etype](feat_dict[srctype]) #apply W_l on src feature
film_weights = self.film[etype](feat_dict[dsttype]) #use dst feature to compute affine function paras
gamma = film_weights[:,:self.out_size] #"gamma" for the affine function
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={}
messages = self.W[etype](
feat_dict[srctype]
) # apply W_l on src feature
film_weights = self.film[etype](
feat_dict[dsttype]
) # use dst feature to compute affine function paras
gamma = film_weights[
:, : self.out_size
] # "gamma" for the affine function
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:
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
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__()
self.film_layers = nn.ModuleList()
self.film_layers.append(
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(
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):
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:
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)
return h
def main(args):
# Step 1: Prepare graph data and retrieve train/validation/test dataloader ============================= #
if args.gpu >= 0 and torch.cuda.is_available():
device = 'cuda:{}'.format(args.gpu)
device = "cuda:{}".format(args.gpu)
else:
device = 'cpu'
device = "cpu"
if args.dataset == 'PPI':
train_set, valid_set, test_set, etypes, in_size, out_size = load_PPI(args.batch_size, device)
if args.dataset == "PPI":
train_set, valid_set, test_set, etypes, in_size, out_size = load_PPI(
args.batch_size, device
)
# 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()
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.step_size, gamma=args.gamma)
optimizer = torch.optim.Adam(
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 ============================================================================== #
lastf1 = 0
......@@ -101,10 +138,12 @@ def main(args):
for batch in train_set:
g = batch.graph
g = g.to(device)
logits = model.forward(g, '_N')
logits = model.forward(g, "_N")
labels = batch.label
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()
loss.backward()
......@@ -112,7 +151,7 @@ def main(args):
train_loss.append(loss.item())
train_f1.append(f1)
train_loss = np.mean(train_loss)
train_loss = np.mean(train_loss)
train_f1 = np.mean(train_f1)
scheduler.step()
......@@ -121,24 +160,32 @@ def main(args):
for batch in valid_set:
g = batch.graph
g = g.to(device)
logits = model.forward(g, '_N')
logits = model.forward(g, "_N")
labels = batch.label
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_f1.append(f1)
val_loss = np.mean(val_loss)
val_loss = np.mean(val_loss)
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:
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:
cnt += 1
if cnt == args.early_stopping:
print('Early stop.')
print("Early stop.")
break
else:
cnt = 0
......@@ -152,41 +199,89 @@ def main(args):
for batch in test_set:
g = batch.graph
g = g.to(device)
logits = model.forward(g, '_N')
logits = model.forward(g, "_N")
labels = batch.label
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_f1.append(f1)
test_loss = np.mean(test_loss)
test_loss = np.mean(test_loss)
test_f1 = np.mean(test_f1)
print("Test F1: {:.4f} | Test loss: {:.4f}".format(test_f1, test_loss))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='GNN-FiLM')
parser.add_argument("--dataset", type=str, default="PPI", help="DGL dataset for this GNN-FiLM")
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")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="GNN-FiLM")
parser.add_argument(
"--dataset",
type=str,
default="PPI",
help="DGL dataset for this GNN-FiLM",
)
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("--max_epoch", type=int, default=1500, help="The max number of epoches. Default: 500")
parser.add_argument("--early_stopping", 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.")
parser.add_argument(
"--max_epoch",
type=int,
default=1500,
help="The max number of epoches. Default: 500",
)
parser.add_argument(
"--early_stopping",
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()
print(args)
if not os.path.exists(args.save_dir):
os.mkdir(args.save_dir)
main(args)
from sklearn.metrics import f1_score
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):
pred = np.round(pred, 0).astype(np.int16)
pred = pred.flatten()
......
......@@ -3,49 +3,100 @@ Data utils for processing bAbI datasets
"""
import os
import string
import torch
from torch.utils.data import DataLoader
import dgl
import torch
import string
from dgl.data.utils import download, get_download_dir, _get_dgl_url, extract_archive
from dgl.data.utils import (
_get_dgl_url,
download,
extract_archive,
get_download_dir,
)
def get_babi_dataloaders(batch_size, train_size=50, task_id=4, q_type=0):
_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:
edge_dict = {'n': 0, 's': 1, 'w': 2, 'e': 3}
edge_dict = {"n": 0, "s": 1, "w": 2, "e": 3}
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:
edge_dict = {'is': 0, 'has_fear': 1}
edge_dict = {"is": 0, "has_fear": 1}
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:
edge_dict = {'is': 0, 'has_color': 1}
edge_dict = {"is": 0, "has_color": 1}
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:
edge_dict = {'>': 0, '<': 1}
label_dict = {'false': 0, 'true': 1}
edge_dict = {">": 0, "<": 1}
label_dict = {"false": 0, "true": 1}
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:
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}
max_seq_length = 2
return _path_finding_dataloader(train_size, batch_size, 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):
return _path_finding_dataloader(
train_size,
batch_size,
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):
graphs = []
labels = []
for d in batch:
edges = d['edges']
edges = d["edges"]
node_ids = []
for s, e, t in edges:
......@@ -55,12 +106,12 @@ def _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse
node_ids.append(t)
g = dgl.DGLGraph()
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)))))
# convert label to node index
label = d['eval'][2]
label = d["eval"][2]
label_idx = nid2idx[label]
labels.append(label_idx)
......@@ -71,19 +122,26 @@ def _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse
if e in reverse_edge:
g.add_edge(nid2idx[t], nid2idx[s])
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[nid2idx[d['eval'][0]]] = 1
g.ndata['annotation'] = annotation.unsqueeze(-1)
annotation[nid2idx[d["eval"][0]]] = 1
g.ndata["annotation"] = annotation.unsqueeze(-1)
graphs.append(g)
batch_graph = dgl.batch(graphs)
labels = torch.tensor(labels, dtype=torch.long)
return batch_graph, labels
def _get_dataloader(data, shuffle):
return DataLoader(dataset=data, 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)
return DataLoader(
dataset=data,
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)
dev_dataloader = _get_dataloader(dev_set, False)
test_dataloaders = []
......@@ -100,26 +158,36 @@ def _convert_ns_dataset(train_size, node_dict, edge_dict, path, q_type):
def convert(file):
dataset = []
d = dict()
with open(file, 'r') as f:
with open(file, "r") as f:
for i, line in enumerate(f.readlines()):
line = line.strip().split()
if line[0] == '1' and len(d) > 0:
if line[0] == "1" and len(d) > 0:
d = dict()
if line[1] == 'eval':
if line[1] == "eval":
# (src, edge, label)
d['eval'] = (node_dict[line[2]], edge_dict[line[3]], node_dict[line[4]])
if d['eval'][1] == q_type:
d["eval"] = (
node_dict[line[2]],
edge_dict[line[3]],
node_dict[line[4]],
)
if d["eval"][1] == q_type:
dataset.append(d)
if len(dataset) >= total_num:
break
else:
if 'edges' not in d:
d['edges'] = []
d['edges'].append((node_dict[line[1]], edge_dict[line[2]], node_dict[line[3]]))
if "edges" not in d:
d["edges"] = []
d["edges"].append(
(
node_dict[line[1]],
edge_dict[line[2]],
node_dict[line[3]],
)
)
return dataset
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)
assert len(data) == total_num
......@@ -128,18 +196,27 @@ def _convert_ns_dataset(train_size, node_dict, edge_dict, path, q_type):
dev_set = data[950:1000]
test_sets = []
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)
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):
graphs = []
labels = []
for d in batch:
edges = d['edges']
edges = d["edges"]
node_ids = []
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
node_ids.append(t)
g = dgl.DGLGraph()
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)))))
labels.append(d['eval'][-1])
labels.append(d["eval"][-1])
edge_types = []
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
if e in reverse_edge:
g.add_edge(nid2idx[t], nid2idx[s])
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[nid2idx[d['eval'][0]]][0] = 1
annotation[nid2idx[d['eval'][2]]][1] = 1
g.ndata['annotation'] = annotation
annotation[nid2idx[d["eval"][0]]][0] = 1
annotation[nid2idx[d["eval"][2]]][1] = 1
g.ndata["annotation"] = annotation
graphs.append(g)
batch_graph = dgl.batch(graphs)
labels = torch.tensor(labels, dtype=torch.long)
return batch_graph, labels
def _get_dataloader(data, shuffle):
return DataLoader(dataset=data, 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)
return DataLoader(
dataset=data,
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)
dev_dataloader = _get_dataloader(dev_set, False)
test_dataloaders = []
......@@ -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
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
def convert(file):
dataset = []
d = dict()
with open(file, 'r') as f:
with open(file, "r") as f:
for i, line in enumerate(f.readlines()):
line = line.strip().split()
if line[0] == '1' and len(d) > 0:
if line[0] == "1" and len(d) > 0:
d = dict()
if line[1] == 'eval':
if line[1] == "eval":
# (src, edge, label)
if 'eval' not in d:
d['eval'] = (node_dict[line[2]], edge_dict[line[3]], node_dict[line[4]], label_dict[line[5]])
if d['eval'][1] == q_type:
if "eval" not in d:
d["eval"] = (
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)
if len(dataset) >= total_num:
break
else:
if 'edges' not in d:
d['edges'] = []
d['edges'].append((node_dict[line[1]], edge_dict[line[2]], node_dict[line[3]]))
if "edges" not in d:
d["edges"] = []
d["edges"].append(
(
node_dict[line[1]],
edge_dict[line[2]],
node_dict[line[3]],
)
)
return dataset
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)
assert len(data) == total_num
......@@ -221,19 +318,27 @@ def _convert_gc_dataset(train_size, node_dict, edge_dict, label_dict, path, q_ty
dev_set = data[950:1000]
test_sets = []
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)
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):
graphs = []
ground_truths = []
seq_lengths = []
for d in batch:
edges = d['edges']
edges = d["edges"]
node_ids = []
for s, e, t in edges:
......@@ -243,12 +348,14 @@ def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, rever
node_ids.append(t)
g = dgl.DGLGraph()
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)))))
truth = d['seq_out'] + [edge_dict['<end>']] * (max_seq_length - len(d['seq_out']))
seq_len = len(d['seq_out'])
truth = d["seq_out"] + [edge_dict["<end>"]] * (
max_seq_length - len(d["seq_out"])
)
seq_len = len(d["seq_out"])
ground_truths.append(truth)
seq_lengths.append(seq_len)
......@@ -259,11 +366,11 @@ def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, rever
if e in reverse_edge:
g.add_edge(nid2idx[t], nid2idx[s])
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[nid2idx[d['eval'][0]]][0] = 1
annotation[nid2idx[d['eval'][1]]][1] = 1
g.ndata['annotation'] = annotation
annotation[nid2idx[d["eval"][0]]][0] = 1
annotation[nid2idx[d["eval"][1]]][1] = 1
g.ndata["annotation"] = annotation
graphs.append(g)
batch_graph = dgl.batch(graphs)
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
return batch_graph, ground_truths, seq_lengths
def _get_dataloader(data, shuffle):
return DataLoader(dataset=data, 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)
return DataLoader(
dataset=data,
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)
dev_dataloader = _get_dataloader(dev_set, False)
test_dataloaders = []
......@@ -290,29 +404,35 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path):
def convert(file):
dataset = []
d = dict()
with open(file, 'r') as f:
with open(file, "r") as f:
for line in f.readlines():
line = line.strip().split()
if line[0] == '1' and len(d) > 0:
if line[0] == "1" and len(d) > 0:
d = dict()
if line[1] == 'eval':
if line[1] == "eval":
# (src, edge, label)
d['eval'] = (node_dict[line[3]], node_dict[line[4]])
d['seq_out'] = []
seq_out = line[5].split(',')
d["eval"] = (node_dict[line[3]], node_dict[line[4]])
d["seq_out"] = []
seq_out = line[5].split(",")
for e in seq_out:
d['seq_out'].append(edge_dict[e])
d["seq_out"].append(edge_dict[e])
dataset.append(d)
if len(dataset) >= total_num:
break
else:
if 'edges' not in d:
d['edges'] = []
d['edges'].append((node_dict[line[1]], edge_dict[line[2]], node_dict[line[3]]))
if "edges" not in d:
d["edges"] = []
d["edges"].append(
(
node_dict[line[1]],
edge_dict[line[2]],
node_dict[line[3]],
)
)
return dataset
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)
assert len(data) == total_num
......@@ -321,7 +441,7 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path):
dev_set = data[950:1000]
test_sets = []
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)
return train_set, dev_set, test_sets
......@@ -329,11 +449,11 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path):
def _download_babi_data():
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)
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):
extract_archive(zip_file_path, extract_dir)
\ No newline at end of file
extract_archive(zip_file_path, extract_dir)
"""
Gated Graph Neural Network module for graph classification tasks
"""
from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
import torch
from torch import nn
from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
class GraphClsGGNN(nn.Module):
def __init__(self,
annotation_size,
out_feats,
n_steps,
n_etypes,
num_cls):
def __init__(self, annotation_size, out_feats, n_steps, n_etypes, num_cls):
super(GraphClsGGNN, self).__init__()
self.annotation_size = annotation_size
self.out_feats = out_feats
self.ggnn = GatedGraphConv(in_feats=out_feats,
out_feats=out_feats,
n_steps=n_steps,
n_etypes=n_etypes)
self.ggnn = GatedGraphConv(
in_feats=out_feats,
out_feats=out_feats,
n_steps=n_steps,
n_etypes=n_etypes,
)
pooling_gate_nn = nn.Linear(annotation_size + out_feats, 1)
self.pooling = GlobalAttentionPooling(pooling_gate_nn)
......@@ -30,16 +28,18 @@ class GraphClsGGNN(nn.Module):
self.loss_fn = nn.CrossEntropyLoss()
def forward(self, graph, labels=None):
etypes = graph.edata.pop('type')
annotation = graph.ndata.pop('annotation').float()
etypes = graph.edata.pop("type")
annotation = graph.ndata.pop("annotation").float()
assert annotation.size()[-1] == self.annotation_size
node_num = graph.number_of_nodes()
zero_pad = torch.zeros([node_num, self.out_feats - self.annotation_size],
dtype=torch.float,
device=annotation.device)
zero_pad = torch.zeros(
[node_num, self.out_feats - self.annotation_size],
dtype=torch.float,
device=annotation.device,
)
h1 = torch.cat([annotation, zero_pad], -1)
out = self.ggnn(graph, h1, etypes)
......@@ -54,4 +54,4 @@ class GraphClsGGNN(nn.Module):
if labels is not None:
loss = self.loss_fn(logits, labels)
return loss, preds
return preds
\ No newline at end of file
return preds
"""
Gated Graph Neural Network module for node selection tasks
"""
from dgl.nn.pytorch import GatedGraphConv
import torch
from torch import nn
import dgl
from dgl.nn.pytorch import GatedGraphConv
class NodeSelectionGGNN(nn.Module):
def __init__(self,
annotation_size,
out_feats,
n_steps,
n_etypes):
def __init__(self, annotation_size, out_feats, n_steps, n_etypes):
super(NodeSelectionGGNN, self).__init__()
self.annotation_size = annotation_size
self.out_feats = out_feats
self.ggnn = GatedGraphConv(in_feats=out_feats,
out_feats=out_feats,
n_steps=n_steps,
n_etypes=n_etypes)
self.ggnn = GatedGraphConv(
in_feats=out_feats,
out_feats=out_feats,
n_steps=n_steps,
n_etypes=n_etypes,
)
self.output_layer = nn.Linear(annotation_size + out_feats, 1)
self.loss_fn = nn.CrossEntropyLoss()
def forward(self, graph, labels=None):
etypes = graph.edata.pop('type')
annotation = graph.ndata.pop('annotation').float()
etypes = graph.edata.pop("type")
annotation = graph.ndata.pop("annotation").float()
assert annotation.size()[-1] == self.annotation_size
node_num = graph.number_of_nodes()
zero_pad = torch.zeros([node_num, self.out_feats - self.annotation_size],
dtype=torch.float,
device=annotation.device)
zero_pad = torch.zeros(
[node_num, self.out_feats - self.annotation_size],
dtype=torch.float,
device=annotation.device,
)
h1 = torch.cat([annotation, zero_pad], -1)
out = self.ggnn(graph, h1, etypes)
all_logits = self.output_layer(torch.cat([out, annotation], -1)).squeeze(-1)
graph.ndata['logits'] = all_logits
all_logits = self.output_layer(
torch.cat([out, annotation], -1)
).squeeze(-1)
graph.ndata["logits"] = all_logits
batch_g = dgl.unbatch(graph)
......@@ -50,7 +53,7 @@ class NodeSelectionGGNN(nn.Module):
if labels is not None:
loss = 0.0
for i, g in enumerate(batch_g):
logits = g.ndata['logits']
logits = g.ndata["logits"]
preds.append(torch.argmax(logits))
if labels is not None:
logits = logits.unsqueeze(0)
......@@ -60,4 +63,4 @@ class NodeSelectionGGNN(nn.Module):
if labels is not None:
loss /= float(len(batch_g))
return loss, preds
return preds
\ No newline at end of file
return preds
......@@ -2,42 +2,49 @@
Gated Graph Sequence Neural Network for sequence outputs
"""
from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
import torch
from torch import nn
import torch.nn.functional as F
from torch import nn
from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
class GGSNN(nn.Module):
def __init__(self,
annotation_size,
out_feats,
n_steps,
n_etypes,
max_seq_length,
num_cls):
def __init__(
self,
annotation_size,
out_feats,
n_steps,
n_etypes,
max_seq_length,
num_cls,
):
super(GGSNN, self).__init__()
self.annotation_size = annotation_size
self.out_feats = out_feats
self.max_seq_length = max_seq_length
self.ggnn = GatedGraphConv(in_feats=out_feats,
out_feats=out_feats,
n_steps=n_steps,
n_etypes=n_etypes)
self.ggnn = GatedGraphConv(
in_feats=out_feats,
out_feats=out_feats,
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)
self.pooling = GlobalAttentionPooling(pooling_gate_nn)
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):
etypes = graph.edata.pop('type')
annotation = graph.ndata.pop('annotation').float()
etypes = graph.edata.pop("type")
annotation = graph.ndata.pop("annotation").float()
assert annotation.size()[-1] == self.annotation_size
......@@ -45,9 +52,11 @@ class GGSNN(nn.Module):
all_logits = []
for _ in range(self.max_seq_length):
zero_pad = torch.zeros([node_num, self.out_feats - self.annotation_size],
dtype=torch.float,
device=annotation.device)
zero_pad = torch.zeros(
[node_num, self.out_feats - self.annotation_size],
dtype=torch.float,
device=annotation.device,
)
h1 = torch.cat([annotation.detach(), zero_pad], -1)
out = self.ggnn(graph, h1, etypes)
......@@ -71,14 +80,18 @@ def sequence_loss(logits, ground_truth, seq_length=None):
def sequence_mask(length):
max_length = logits.size(1)
batch_size = logits.size(0)
range_tensor = torch.arange(0, max_length, dtype=seq_length.dtype, device=seq_length.device)
range_tensor = torch.stack([range_tensor]*batch_size, 0)
range_tensor = torch.arange(
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()
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:
loss = loss.mean()
......@@ -86,4 +99,4 @@ def sequence_loss(logits, ground_truth, seq_length=None):
mask = sequence_mask(seq_length)
loss = (loss * mask).sum(-1) / seq_length.float()
loss = loss.mean()
return loss
\ No newline at end of file
return loss
......@@ -3,34 +3,38 @@ Training and testing for graph classification tasks in bAbI
"""
import argparse
import numpy as np
import torch
from data_utils import get_babi_dataloaders
from ggnn_gc import GraphClsGGNN
from torch.optim import Adam
import torch
import numpy as np
def main(args):
out_feats = {18: 3}
n_etypes = {18: 2}
train_dataloader, dev_dataloader, test_dataloaders = \
get_babi_dataloaders(batch_size=args.batch_size,
train_size=args.train_num,
task_id=args.task_id,
q_type=args.question_id)
model = GraphClsGGNN(annotation_size=2,
out_feats=out_feats[args.task_id],
n_steps=5,
n_etypes=n_etypes[args.task_id],
num_cls=2)
train_dataloader, dev_dataloader, test_dataloaders = get_babi_dataloaders(
batch_size=args.batch_size,
train_size=args.train_num,
task_id=args.task_id,
q_type=args.question_id,
)
model = GraphClsGGNN(
annotation_size=2,
out_feats=out_feats[args.task_id],
n_steps=5,
n_etypes=n_etypes[args.task_id],
num_cls=2,
)
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'Dev set size: {len(dev_dataloader.dataset)}')
print(f"Training set size: {len(train_dataloader.dataset)}")
print(f"Dev set size: {len(dev_dataloader.dataset)}")
# training and dev stage
for epoch in range(args.epochs):
......@@ -42,7 +46,7 @@ def main(args):
loss.backward()
opt.step()
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:
continue
......@@ -62,7 +66,7 @@ def main(args):
# test stage
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 = []
for dataloader in test_dataloaders:
......@@ -83,24 +87,30 @@ def main(args):
test_acc_mean = np.mean(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}')
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,
help='task id from 1 to 20')
parser.add_argument('--question_id', type=int, default=0,
help='question id for each task')
parser.add_argument('--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')
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, help="task id from 1 to 20"
)
parser.add_argument(
"--question_id", type=int, default=0, help="question id for each task"
)
parser.add_argument(
"--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()
main(args)
\ No newline at end of file
main(args)
......@@ -3,34 +3,38 @@ Training and testing for node selection tasks in bAbI
"""
import argparse
import time
import numpy as np
import torch
from data_utils import get_babi_dataloaders
from ggnn_ns import NodeSelectionGGNN
from torch.optim import Adam
import torch
import numpy as np
import time
def main(args):
out_feats = {4: 4, 15: 5, 16: 6}
n_etypes = {4: 4, 15: 2, 16: 2}
train_dataloader, dev_dataloader, test_dataloaders = \
get_babi_dataloaders(batch_size=args.batch_size,
train_size=args.train_num,
task_id=args.task_id,
q_type=args.question_id)
model = NodeSelectionGGNN(annotation_size=1,
out_feats=out_feats[args.task_id],
n_steps=5,
n_etypes=n_etypes[args.task_id])
train_dataloader, dev_dataloader, test_dataloaders = get_babi_dataloaders(
batch_size=args.batch_size,
train_size=args.train_num,
task_id=args.task_id,
q_type=args.question_id,
)
model = NodeSelectionGGNN(
annotation_size=1,
out_feats=out_feats[args.task_id],
n_steps=5,
n_etypes=n_etypes[args.task_id],
)
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'Dev set size: {len(dev_dataloader.dataset)}')
print(f"Training set size: {len(train_dataloader.dataset)}")
print(f"Dev set size: {len(dev_dataloader.dataset)}")
# training and dev stage
for epoch in range(args.epochs):
......@@ -41,7 +45,7 @@ def main(args):
opt.zero_grad()
loss.backward()
opt.step()
print(f'Epoch {epoch}, batch {i} loss: {loss.data}')
print(f"Epoch {epoch}, batch {i} loss: {loss.data}")
dev_preds = []
dev_labels = []
......@@ -49,7 +53,9 @@ def main(args):
for g, labels in dev_dataloader:
with torch.no_grad():
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()
dev_preds += preds
dev_labels += labels
......@@ -59,7 +65,7 @@ def main(args):
# test stage
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 = []
for dataloader in test_dataloaders:
......@@ -69,7 +75,9 @@ def main(args):
for g, labels in dataloader:
with torch.no_grad():
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()
test_preds += preds
test_labels += labels
......@@ -80,24 +88,30 @@ def main(args):
test_acc_mean = np.mean(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}')
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,
help='task id from 1 to 20')
parser.add_argument('--question_id', type=int, default=1,
help='question id for each task')
parser.add_argument('--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')
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, help="task id from 1 to 20"
)
parser.add_argument(
"--question_id", type=int, default=1, help="question id for each task"
)
parser.add_argument(
"--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()
main(args)
\ No newline at end of file
main(args)
......@@ -4,35 +4,39 @@ Here we take task 19 'Path Finding' as an example
"""
import argparse
import numpy as np
import torch
from data_utils import get_babi_dataloaders
from ggsnn import GGSNN
from torch.optim import Adam
import torch
import numpy as np
def main(args):
out_feats = {19: 6}
n_etypes = {19: 4}
train_dataloader, dev_dataloader, test_dataloaders = \
get_babi_dataloaders(batch_size=args.batch_size,
train_size=args.train_num,
task_id=args.task_id,
q_type=-1)
model = GGSNN(annotation_size=2,
out_feats=out_feats[args.task_id],
n_steps=5,
n_etypes=n_etypes[args.task_id],
max_seq_length=2,
num_cls=5)
train_dataloader, dev_dataloader, test_dataloaders = get_babi_dataloaders(
batch_size=args.batch_size,
train_size=args.train_num,
task_id=args.task_id,
q_type=-1,
)
model = GGSNN(
annotation_size=2,
out_feats=out_feats[args.task_id],
n_steps=5,
n_etypes=n_etypes[args.task_id],
max_seq_length=2,
num_cls=5,
)
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'Dev set size: {len(dev_dataloader.dataset)}')
print(f"Training set size: {len(train_dataloader.dataset)}")
print(f"Dev set size: {len(dev_dataloader.dataset)}")
# training and dev stage
for epoch in range(args.epochs):
......@@ -44,7 +48,7 @@ def main(args):
loss.backward()
opt.step()
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:
continue
......@@ -65,7 +69,7 @@ def main(args):
# test stage
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 = []
for dataloader in test_dataloaders:
......@@ -87,23 +91,28 @@ def main(args):
test_acc_mean = np.mean(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}')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Gated Graph Sequence Neural Networks for sequential output tasks in '
'bAbI')
parser.add_argument('--task_id', type=int, default=19,
help='task id from 1 to 20')
parser.add_argument('--train_num', type=int, default=250,
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=200,
help='number of training epochs')
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"
)
parser.add_argument(
"--task_id", type=int, default=19, help="task id from 1 to 20"
)
parser.add_argument(
"--train_num", type=int, default=250, 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=200, help="number of training epochs"
)
args = parser.parse_args()
main(args)
\ No newline at end of file
main(args)
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
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 torch.utils.data.sampler import SubsetRandomSampler
from dgl.data import GINDataset
from dgl.dataloading import GraphDataLoader
from dgl.nn.pytorch.conv import GINConv
from dgl.nn.pytorch.glob import SumPooling
import argparse
class MLP(nn.Module):
"""Construct two-layer MLP-type aggreator for GIN model"""
def __init__(self, input_dim, hidden_dim, output_dim):
super().__init__()
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(hidden_dim, output_dim, bias=False))
self.batch_norm = nn.BatchNorm1d((hidden_dim))
......@@ -25,7 +29,8 @@ class MLP(nn.Module):
h = x
h = F.relu(self.batch_norm(self.linears[0](h)))
return self.linears[1](h)
class GIN(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super().__init__()
......@@ -33,12 +38,14 @@ class GIN(nn.Module):
self.batch_norms = nn.ModuleList()
num_layers = 5
# 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:
mlp = MLP(input_dim, hidden_dim, hidden_dim)
else:
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))
# linear functions for graph sum poolings of output of each layer
self.linear_prediction = nn.ModuleList()
......@@ -48,7 +55,9 @@ class GIN(nn.Module):
else:
self.linear_prediction.append(nn.Linear(hidden_dim, output_dim))
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):
# list of hidden representation at each layer (including the input layer)
......@@ -64,7 +73,8 @@ class GIN(nn.Module):
pooled_h = self.pool(g, h)
score_over_layer += self.drop(self.linear_prediction[i](pooled_h))
return score_over_layer
def split_fold10(labels, fold_idx=0):
skf = StratifiedKFold(n_splits=10, shuffle=True, random_state=0)
idx_list = []
......@@ -73,6 +83,7 @@ def split_fold10(labels, fold_idx=0):
train_idx, valid_idx = idx_list[fold_idx]
return train_idx, valid_idx
def evaluate(dataloader, device, model):
model.eval()
total = 0
......@@ -80,7 +91,7 @@ def evaluate(dataloader, device, model):
for batched_graph, labels in dataloader:
batched_graph = batched_graph.to(device)
labels = labels.to(device)
feat = batched_graph.ndata.pop('attr')
feat = batched_graph.ndata.pop("attr")
total += len(labels)
logits = model(batched_graph, feat)
_, predicted = torch.max(logits, 1)
......@@ -88,20 +99,21 @@ def evaluate(dataloader, device, model):
acc = 1.0 * total_correct / total
return acc
def train(train_loader, val_loader, device, model):
# loss function, optimizer and scheduler
loss_fcn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5)
# training loop
# training loop
for epoch in range(350):
model.train()
total_loss = 0
for batch, (batched_graph, labels) in enumerate(train_loader):
batched_graph = batched_graph.to(device)
labels = labels.to(device)
feat = batched_graph.ndata.pop('attr')
feat = batched_graph.ndata.pop("attr")
logits = model(batched_graph, feat)
loss = loss_fcn(logits, labels)
optimizer.zero_grad()
......@@ -111,34 +123,52 @@ def train(train_loader, val_loader, device, model):
scheduler.step()
train_acc = evaluate(train_loader, device, model)
valid_acc = evaluate(val_loader, device, model)
print("Epoch {:05d} | Loss {:.4f} | Train Acc. {:.4f} | Validation Acc. {:.4f} "
. format(epoch, total_loss / (batch + 1), train_acc, valid_acc))
print(
"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.add_argument('--dataset', type=str, default="MUTAG",
choices=['MUTAG', 'PTC', 'NCI1', 'PROTEINS'],
help='name of dataset (default: MUTAG)')
parser.add_argument(
"--dataset",
type=str,
default="MUTAG",
choices=["MUTAG", "PTC", "NCI1", "PROTEINS"],
help="name of dataset (default: MUTAG)",
)
args = parser.parse_args()
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')
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")
# 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]
train_idx, val_idx = split_fold10(labels)
# create dataloader
train_loader = GraphDataLoader(dataset, sampler=SubsetRandomSampler(train_idx),
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())
train_loader = GraphDataLoader(
dataset,
sampler=SubsetRandomSampler(train_idx),
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
in_size = dataset.dim_nfeats
out_size = dataset.gclasses
model = GIN(in_size, 16, out_size).to(device)
# model training/validating
print('Training...')
print("Training...")
train(train_loader, val_loader, device, model)
# Data augmentation on graphs via edge dropping and feature masking
import torch as th
import numpy as np
import torch as th
import dgl
def aug(graph, x, feat_drop_rate, edge_mask_rate):
n_node = graph.num_nodes()
......@@ -21,19 +23,22 @@ def aug(graph, x, feat_drop_rate, edge_mask_rate):
return ng, feat
def drop_feature(x, drop_prob):
drop_mask = th.empty((x.size(1),),
dtype=th.float32,
device=x.device).uniform_(0, 1) < drop_prob
drop_mask = (
th.empty((x.size(1),), dtype=th.float32, device=x.device).uniform_(0, 1)
< drop_prob
)
x = x.clone()
x[:, drop_mask] = 0
return x
def mask_edge(graph, mask_prob):
E = graph.num_edges()
mask_rates = th.FloatTensor(np.ones(E) * mask_prob)
masks = th.bernoulli(1 - mask_rates)
mask_idx = masks.nonzero().squeeze(1)
return mask_idx
\ No newline at end of file
return mask_idx
from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
def load(name):
if name == 'cora':
if name == "cora":
dataset = CoraGraphDataset()
elif name == 'citeseer':
elif name == "citeseer":
dataset = CiteseerGraphDataset()
elif name == 'pubmed':
elif name == "pubmed":
dataset = PubmedGraphDataset()
graph = dataset[0]
train_mask = graph.ndata.pop('train_mask')
test_mask = graph.ndata.pop('test_mask')
train_mask = graph.ndata.pop("train_mask")
test_mask = graph.ndata.pop("test_mask")
feat = graph.ndata.pop('feat')
labels = graph.ndata.pop('label')
feat = graph.ndata.pop("feat")
labels = graph.ndata.pop("label")
return graph, feat, labels, train_mask, test_mask
\ No newline at end of file
return graph, feat, labels, train_mask, test_mask
'''
"""
Code adapted from https://github.com/CRIPAC-DIG/GRACE
Linear evaluation on learned node embeddings
'''
"""
import numpy as np
import functools
from sklearn.metrics import f1_score
import numpy as np
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.preprocessing import normalize, OneHotEncoder
from sklearn.preprocessing import OneHotEncoder, normalize
def repeat(n_times):
......@@ -22,8 +22,9 @@ def repeat(n_times):
for key in results[0].keys():
values = [r[key] for r in results]
statistics[key] = {
'mean': np.mean(values),
'std': np.std(values)}
"mean": np.mean(values),
"std": np.std(values),
}
print_statistics(statistics, f.__name__)
return statistics
......@@ -41,41 +42,49 @@ def prob_to_one_hot(y_pred):
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()):
mean = statistics[key]['mean']
std = statistics[key]['std']
print(f'{key}={mean:.4f}+-{std:.4f}', end='')
mean = statistics[key]["mean"]
std = statistics[key]["std"]
print(f"{key}={mean:.4f}+-{std:.4f}", end="")
if i != len(statistics.keys()) - 1:
print(',', end=' ')
print(",", end=" ")
else:
print()
@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()
Y = y.detach().cpu().numpy()
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)
X = normalize(X, norm='l2')
X = normalize(X, norm="l2")
if split == 'random':
X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=1 - ratio)
elif split == 'public':
if split == "random":
X_train, X_test, y_train, y_test = train_test_split(
X, Y, test_size=1 - ratio
)
elif split == "public":
X_train = X[train_mask]
X_test = X[test_mask]
y_train = Y[train_mask]
y_test = Y[test_mask]
logreg = LogisticRegression(solver='liblinear')
logreg = LogisticRegression(solver="liblinear")
c = 2.0 ** np.arange(-10, 10)
clf = GridSearchCV(estimator=OneVsRestClassifier(logreg),
param_grid=dict(estimator__C=c), n_jobs=8, cv=5,
verbose=0)
clf = GridSearchCV(
estimator=OneVsRestClassifier(logreg),
param_grid=dict(estimator__C=c),
n_jobs=8,
cv=5,
verbose=0,
)
clf.fit(X_train, y_train)
y_pred = clf.predict_proba(X_test)
......@@ -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")
macro = f1_score(y_test, y_pred, average="macro")
return {
'F1Mi': micro,
'F1Ma': macro
}
return {"F1Mi": micro, "F1Ma": macro}
import argparse
from model import Grace
from aug import aug
from dataset import load
import warnings
import numpy as np
import torch as th
import torch.nn as nn
from aug import aug
from dataset import load
from eval import label_classification
import warnings
from model import Grace
warnings.filterwarnings('ignore')
warnings.filterwarnings("ignore")
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.add_argument('--dataname', type=str, default='cora')
parser.add_argument('--gpu', type=int, default=0)
parser.add_argument('--split', type=str, default='random')
parser.add_argument('--epochs', type=int, default=500, help='Number of training periods.')
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('--temp', type=float, default=1.0, help='Temperature.')
parser.add_argument('--act_fn', type=str, default='relu')
parser.add_argument("--hid_dim", type=int, default=256, help='Hidden layer dim.')
parser.add_argument("--out_dim", type=int, default=256, help='Output layer dim.')
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.')
parser.add_argument("--dataname", type=str, default="cora")
parser.add_argument("--gpu", type=int, default=0)
parser.add_argument("--split", type=str, default="random")
parser.add_argument(
"--epochs", type=int, default=500, help="Number of training periods."
)
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("--temp", type=float, default=1.0, help="Temperature.")
parser.add_argument("--act_fn", type=str, default="relu")
parser.add_argument(
"--hid_dim", type=int, default=256, help="Hidden layer dim."
)
parser.add_argument(
"--out_dim", type=int, default=256, help="Output layer dim."
)
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()
if args.gpu != -1 and th.cuda.is_available():
args.device = 'cuda:{}'.format(args.gpu)
args.device = "cuda:{}".format(args.gpu)
else:
args.device = 'cpu'
args.device = "cpu"
if __name__ == '__main__':
if __name__ == "__main__":
# Step 1: Load hyperparameters =================================================================== #
lr = args.lr
......@@ -53,7 +82,7 @@ if __name__ == '__main__':
out_dim = args.out_dim
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_2 = args.der2
......@@ -71,7 +100,7 @@ if __name__ == '__main__':
# Step 3: Create model =================================================================== #
model = Grace(in_dim, hid_dim, out_dim, num_layers, act_fn, temp)
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)
......@@ -92,7 +121,7 @@ if __name__ == '__main__':
loss.backward()
optimizer.step()
print(f'Epoch={epoch:03d}, loss={loss.item():.4f}')
print(f"Epoch={epoch:03d}, loss={loss.item():.4f}")
# Step 5: Linear evaluation ============================================================== #
print("=== Final ===")
......@@ -102,5 +131,7 @@ if __name__ == '__main__':
feat = feat.to(args.device)
embeds = model.get_embedding(graph, feat)
'''Evaluation Embeddings '''
label_classification(embeds, labels, train_mask, test_mask, split=args.split)
\ No newline at end of file
"""Evaluation Embeddings """
label_classification(
embeds, labels, train_mask, test_mask, split=args.split
)
import torch as th
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn import GraphConv
# Multi-layer Graph Convolutional Networks
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__()
assert num_layers >= 2
......@@ -26,6 +27,7 @@ class GCN(nn.Module):
return feat
# Multi-layer(2-layer) Perceptron
class MLP(nn.Module):
def __init__(self, in_dim, out_dim):
......@@ -56,6 +58,7 @@ class Grace(nn.Module):
temp: float
Temperature constant.
"""
def __init__(self, in_dim, hid_dim, out_dim, num_layers, act_fn, temp):
super(Grace, self).__init__()
self.encoder = GCN(in_dim, hid_dim, act_fn, num_layers)
......@@ -74,8 +77,8 @@ class Grace(nn.Module):
# calculate SimCLR loss
f = lambda x: th.exp(x / self.temp)
refl_sim = f(self.sim(z1, z1)) # intra-view pairs
between_sim = f(self.sim(z1, z2)) # inter-view pairs
refl_sim = f(self.sim(z1, z1)) # intra-view pairs
between_sim = f(self.sim(z1, z2)) # inter-view pairs
# between_sim.diag(): positive pairs
x1 = refl_sim.sum(1) + between_sim.sum(1) - refl_sim.diag()
......@@ -104,4 +107,4 @@ class Grace(nn.Module):
ret = (l1 + l2) * 0.5
return ret.mean()
\ No newline at end of file
return ret.mean()
import argparse
import warnings
import numpy as np
import torch as th
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from model import GRAND
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():
parser = argparse.ArgumentParser(description='GRAND')
parser = argparse.ArgumentParser(description="GRAND")
# 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
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
parser.add_argument('--epochs', type=int, default=200, help='Training epochs.')
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.')
parser.add_argument(
"--epochs", type=int, default=200, help="Training epochs."
)
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
parser.add_argument("--hid_dim", type=int, default=32, help='Hidden layer dimensionalities.')
parser.add_argument('--dropnode_rate', type=float, default=0.5,
help='Dropnode rate (1 - keep probability).')
parser.add_argument('--input_droprate', 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., help='Coefficient of consistency regularization')
parser.add_argument('--use_bn', action='store_true', default=False, help='Using Batch Normalization')
parser.add_argument(
"--hid_dim", type=int, default=32, help="Hidden layer dimensionalities."
)
parser.add_argument(
"--dropnode_rate",
type=float,
default=0.5,
help="Dropnode rate (1 - keep probability).",
)
parser.add_argument(
"--input_droprate",
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()
# check cuda
if args.gpu != -1 and th.cuda.is_available():
args.device = 'cuda:{}'.format(args.gpu)
args.device = "cuda:{}".format(args.gpu)
else:
args.device = 'cpu'
args.device = "cpu"
return args
def consis_loss(logps, temp, lam):
ps = [th.exp(p) for p in logps]
ps = th.stack(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()
ps = th.stack(ps, dim=2)
avg_p = th.mean(ps, dim=2)
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)
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
return loss
if __name__ == '__main__':
if __name__ == "__main__":
# Step 1: Prepare graph data and retrieve train/validation/test index ============================= #
# Load from DGL dataset
args = argument()
print(args)
if args.dataname == 'cora':
if args.dataname == "cora":
dataset = CoraGraphDataset()
elif args.dataname == 'citeseer':
elif args.dataname == "citeseer":
dataset = CiteseerGraphDataset()
elif args.dataname == 'pubmed':
elif args.dataname == "pubmed":
dataset = PubmedGraphDataset()
graph = dataset[0]
graph = dgl.add_self_loop(graph)
device = args.device
......@@ -86,100 +133,121 @@ if __name__ == '__main__':
n_classes = dataset.num_classes
# retrieve labels of ground truth
labels = graph.ndata.pop('label').to(device).long()
labels = graph.ndata.pop("label").to(device).long()
# Extract node features
feats = graph.ndata.pop('feat').to(device)
feats = graph.ndata.pop("feat").to(device)
n_features = feats.shape[-1]
# retrieve masks for train/validation/test
train_mask = graph.ndata.pop('train_mask')
val_mask = graph.ndata.pop('val_mask')
test_mask = graph.ndata.pop('test_mask')
train_mask = graph.ndata.pop("train_mask")
val_mask = graph.ndata.pop("val_mask")
test_mask = graph.ndata.pop("test_mask")
train_idx = th.nonzero(train_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)
# Step 2: Create model =================================================================== #
model = GRAND(n_features, args.hid_dim, n_classes, args.sample, args.order,
args.dropnode_rate, args.input_droprate,
args.hidden_droprate, args.use_bn)
model = GRAND(
n_features,
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)
graph = graph.to(args.device)
# Step 3: Create training components ===================================================== #
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
acc_best = 0
# Step 4: training epoches =============================================================== #
for epoch in range(args.epochs):
''' Training '''
"""Training"""
model.train()
loss_sup = 0
logits = model(graph, feats, True)
# calculate supervised loss
for k in range(args.sample):
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
loss_consis = consis_loss(logits, args.tem, args.lam)
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
opt.zero_grad()
loss_train.backward()
opt.step()
''' Validating '''
""" Validating """
model.eval()
with th.no_grad():
val_logits = model(graph, feats, False)
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)
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)
# Print out performance
print("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()))
print(
"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
if loss_val < loss_best or acc_val > acc_best:
if loss_val < loss_best:
best_epoch = epoch
th.save(model.state_dict(), args.dataname +'.pkl')
th.save(model.state_dict(), args.dataname + ".pkl")
no_improvement = 0
loss_best = min(loss_val, loss_best)
acc_best = max(acc_val, acc_best)
else:
no_improvement += 1
if no_improvement == args.early_stopping:
print('Early stopping.')
print("Early stopping.")
break
print("Optimization Finished!")
print('Loading {}th epoch'.format(best_epoch))
model.load_state_dict(th.load(args.dataname +'.pkl'))
''' Testing '''
print("Loading {}th epoch".format(best_epoch))
model.load_state_dict(th.load(args.dataname + ".pkl"))
""" Testing """
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 torch as th
import torch.nn as nn
import dgl.function as fn
import torch.nn.functional as F
import dgl.function as fn
def drop_node(feats, drop_rate, training):
n = feats.shape[0]
drop_rates = th.FloatTensor(np.ones(n) * drop_rate)
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
else:
feats = feats * (1. - drop_rate)
feats = feats * (1.0 - drop_rate)
return feats
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__()
self.layer1 = nn.Linear(nfeat, nhid, bias = True)
self.layer2 = nn.Linear(nhid, nclass, bias = True)
self.layer1 = nn.Linear(nfeat, nhid, bias=True)
self.layer2 = nn.Linear(nhid, nclass, bias=True)
self.input_dropout = nn.Dropout(input_droprate)
self.hidden_dropout = nn.Dropout(hidden_droprate)
self.bn1 = nn.BatchNorm1d(nfeat)
self.bn2 = nn.BatchNorm1d(nhid)
self.use_bn = use_bn
def reset_parameters(self):
self.layer1.reset_parameters()
self.layer2.reset_parameters()
def forward(self, x):
if self.use_bn:
if self.use_bn:
x = self.bn1(x)
x = self.input_dropout(x)
x = F.relu(self.layer1(x))
if self.use_bn:
x = self.bn2(x)
x = self.hidden_dropout(x)
x = self.layer2(x)
return x
return x
def GRANDConv(graph, feats, order):
'''
"""
Parameters
-----------
graph: dgl.Graph
The input graph
feats: Tensor (n_nodes * feat_dim)
Node features
order: int
order: int
Propagation Steps
'''
"""
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)
norm = th.pow(degs, -0.5).to(feats.device).unsqueeze(1)
graph.ndata['norm'] = norm
graph.apply_edges(fn.u_mul_v('norm', 'norm', 'weight'))
''' Graph Conv '''
graph.ndata["norm"] = norm
graph.apply_edges(fn.u_mul_v("norm", "norm", "weight"))
""" Graph Conv """
x = feats
y = 0+feats
y = 0 + feats
for i in range(order):
graph.ndata['h'] = x
graph.update_all(fn.u_mul_e('h', 'weight', 'm'), fn.sum('m', 'h'))
x = graph.ndata.pop('h')
graph.ndata["h"] = x
graph.update_all(fn.u_mul_e("h", "weight", "m"), fn.sum("m", "h"))
x = graph.ndata.pop("h")
y.add_(x)
return y /(order + 1)
return y / (order + 1)
class GRAND(nn.Module):
r"""
......@@ -108,16 +114,19 @@ class GRAND(nn.Module):
If True, use batch normalization.
"""
def __init__(self,
in_dim,
hid_dim,
n_class,
S = 1,
K = 3,
node_dropout=0.0,
input_droprate = 0.0,
hidden_droprate = 0.0,
batchnorm=False):
def __init__(
self,
in_dim,
hid_dim,
n_class,
S=1,
K=3,
node_dropout=0.0,
input_droprate=0.0,
hidden_droprate=0.0,
batchnorm=False,
):
super(GRAND, self).__init__()
self.in_dim = in_dim
......@@ -125,29 +134,31 @@ class GRAND(nn.Module):
self.S = S
self.K = K
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.node_dropout = nn.Dropout(node_dropout)
def forward(self, graph, feats, training = True):
def forward(self, graph, feats, training=True):
X = feats
S = self.S
if training: # Training Mode
if training: # Training Mode
output_list = []
for s in range(S):
drop_feat = drop_node(X, self.dropout, True) # Drop node
feat = GRANDConv(graph, drop_feat, self.K) # Graph Convolution
output_list.append(th.log_softmax(self.mlp(feat), dim=-1)) # Prediction
feat = GRANDConv(graph, drop_feat, self.K) # Graph Convolution
output_list.append(
th.log_softmax(self.mlp(feat), dim=-1)
) # Prediction
return output_list
else: # Inference Mode
drop_feat = drop_node(X, self.dropout, False)
X = GRANDConv(graph, drop_feat, self.K)
else: # Inference Mode
drop_feat = drop_node(X, self.dropout, False)
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
from ged import graph_edit_distance
import dgl
src1 = [0, 1, 2, 3, 4, 5];
dst1 = [1, 2, 3, 4, 5, 6];
src1 = [0, 1, 2, 3, 4, 5]
dst1 = [1, 2, 3, 4, 5, 6]
src2 = [0, 1, 3, 4, 5];
dst2 = [1, 2, 4, 5, 6];
src2 = [0, 1, 3, 4, 5]
dst2 = [1, 2, 4, 5, 6]
G1 = dgl.DGLGraph((src1, dst1))
......@@ -15,55 +15,73 @@ G2 = dgl.DGLGraph((src2, dst2))
# Exact edit distance with astar search
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G1, algorithm='astar')
print(distance) # 0.0
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='astar')
print(distance) # 1.0
distance, node_mapping, edge_mapping = graph_edit_distance(
G1, G1, algorithm="astar"
)
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
node_substitution_cost = np.empty((G1.number_of_nodes(), G2.number_of_nodes()));
G1_node_deletion_cost = np.empty(G1.number_of_nodes());
G2_node_insertion_cost = np.empty(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())
G2_node_insertion_cost = np.empty(G2.number_of_nodes())
edge_substitution_cost = np.empty((G1.number_of_edges(), G2.number_of_edges()));
G1_edge_deletion_cost = np.empty(G1.number_of_edges());
G2_edge_insertion_cost = np.empty(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())
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.fill(1.0);
node_substitution_cost.fill(1.0)
for i in range(G1.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
G1_node_deletion_cost.fill(1.0);
G2_node_insertion_cost.fill(1.0);
G1_node_deletion_cost.fill(1.0)
G2_node_insertion_cost.fill(1.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
G1_edge_deletion_cost.fill(0.5);
G2_edge_insertion_cost.fill(0.5);
G1_edge_deletion_cost.fill(0.5)
G2_edge_insertion_cost.fill(0.5)
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, \
node_substitution_cost, edge_substitution_cost, \
G1_node_deletion_cost, G2_node_insertion_cost, \
G1_edge_deletion_cost, G2_edge_insertion_cost, \
algorithm="astar")
distance, node_mapping, edge_mapping = graph_edit_distance(
G1,
G2,
node_substitution_cost,
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
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='beam', max_beam_size=2)
print(distance) # 3.0
distance, node_mapping, edge_mapping = graph_edit_distance(
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
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='bipartite')
print(distance) # 9.0, can be different as multiple solutions possible for the intermediate LAP used in this approximation
distance, node_mapping, edge_mapping = graph_edit_distance(
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
distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='hausdorff')
print(distance) # 0.0
\ No newline at end of file
distance, node_mapping, edge_mapping = graph_edit_distance(
G1, G2, algorithm="hausdorff"
)
print(distance) # 0.0
import sklearn.linear_model as lm
import sklearn.metrics as skm
import torch as th
import torch.nn as nn
import torch.functional as F
import torch.nn as nn
import tqdm
import dgl
import dgl.nn as dglnn
import sklearn.linear_model as lm
import sklearn.metrics as skm
import tqdm
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__()
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_hidden = n_hidden
self.n_classes = n_classes
self.layers = nn.ModuleList()
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):
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_hidden, "mean"))
self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, "mean"))
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.activation = activation
......@@ -51,7 +57,10 @@ class SAGE(nn.Module):
# on each layer are of course splitted in batches.
# TODO: can we standardize this?
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)
dataloader = dgl.dataloading.DataLoader(
......@@ -62,7 +71,8 @@ class SAGE(nn.Module):
batch_size=batch_size,
shuffle=False,
drop_last=False,
num_workers=num_workers)
num_workers=num_workers,
)
for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
block = blocks[0]
......@@ -79,6 +89,7 @@ class SAGE(nn.Module):
x = y
return y
def compute_acc_unsupervised(emb, labels, train_nids, val_nids, test_nids):
"""
Compute the accuracy of prediction given the labels.
......@@ -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)
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)
pred = lr.predict(emb)
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_eval = skm.f1_score(val_labels, pred[val_nids], average="micro")
f1_micro_test = skm.f1_score(test_labels, pred[test_nids], average="micro")
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