chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,145 @@
|
||||
""" Code adapted from https://github.com/kavehhassani/mvgrl """
|
||||
import dgl
|
||||
import networkx as nx
|
||||
import numpy as np
|
||||
import scipy.sparse as sp
|
||||
import torch as th
|
||||
from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
|
||||
from dgl.nn import APPNPConv
|
||||
from scipy.linalg import fractional_matrix_power, inv
|
||||
from sklearn.preprocessing import MinMaxScaler
|
||||
|
||||
|
||||
def preprocess_features(features):
|
||||
"""Row-normalize feature matrix and convert to tuple representation"""
|
||||
rowsum = np.array(features.sum(1))
|
||||
r_inv = np.power(rowsum, -1).flatten()
|
||||
r_inv[np.isinf(r_inv)] = 0.0
|
||||
r_mat_inv = sp.diags(r_inv)
|
||||
features = r_mat_inv.dot(features)
|
||||
if isinstance(features, np.ndarray):
|
||||
return features
|
||||
else:
|
||||
return features.todense(), sparse_to_tuple(features)
|
||||
|
||||
|
||||
def sparse_to_tuple(sparse_mx):
|
||||
"""Convert sparse matrix to tuple representation."""
|
||||
|
||||
def to_tuple(mx):
|
||||
if not sp.isspmatrix_coo(mx):
|
||||
mx = mx.tocoo()
|
||||
coords = np.vstack((mx.row, mx.col)).transpose()
|
||||
values = mx.data
|
||||
shape = mx.shape
|
||||
return coords, values, shape
|
||||
|
||||
if isinstance(sparse_mx, list):
|
||||
for i in range(len(sparse_mx)):
|
||||
sparse_mx[i] = to_tuple(sparse_mx[i])
|
||||
else:
|
||||
sparse_mx = to_tuple(sparse_mx)
|
||||
|
||||
return sparse_mx
|
||||
|
||||
|
||||
def compute_ppr(graph: nx.Graph, alpha=0.2, self_loop=True):
|
||||
a = nx.convert_matrix.to_numpy_array(graph)
|
||||
if self_loop:
|
||||
a = a + np.eye(a.shape[0]) # A^ = A + I_n
|
||||
d = np.diag(np.sum(a, 1)) # D^ = Sigma A^_ii
|
||||
dinv = fractional_matrix_power(d, -0.5) # D^(-1/2)
|
||||
at = np.matmul(np.matmul(dinv, a), dinv) # A~ = D^(-1/2) x A^ x D^(-1/2)
|
||||
return alpha * inv(
|
||||
(np.eye(a.shape[0]) - (1 - alpha) * at)
|
||||
) # a(I_n-(1-a)A~)^-1
|
||||
|
||||
|
||||
def process_dataset(name, epsilon):
|
||||
if name == "cora":
|
||||
dataset = CoraGraphDataset()
|
||||
elif name == "citeseer":
|
||||
dataset = CiteseerGraphDataset()
|
||||
|
||||
graph = dataset[0]
|
||||
feat = graph.ndata.pop("feat")
|
||||
label = graph.ndata.pop("label")
|
||||
|
||||
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()
|
||||
val_idx = th.nonzero(val_mask, as_tuple=False).squeeze()
|
||||
test_idx = th.nonzero(test_mask, as_tuple=False).squeeze()
|
||||
|
||||
nx_g = dgl.to_networkx(graph)
|
||||
|
||||
print("computing ppr")
|
||||
diff_adj = compute_ppr(nx_g, 0.2)
|
||||
print("computing end")
|
||||
|
||||
if name == "citeseer":
|
||||
print("additional processing")
|
||||
feat = th.tensor(preprocess_features(feat.numpy())).float()
|
||||
diff_adj[diff_adj < epsilon] = 0
|
||||
scaler = MinMaxScaler()
|
||||
scaler.fit(diff_adj)
|
||||
diff_adj = scaler.transform(diff_adj)
|
||||
|
||||
diff_edges = np.nonzero(diff_adj)
|
||||
diff_weight = diff_adj[diff_edges]
|
||||
diff_graph = dgl.graph(diff_edges)
|
||||
|
||||
graph = graph.add_self_loop()
|
||||
|
||||
return (
|
||||
graph,
|
||||
diff_graph,
|
||||
feat,
|
||||
label,
|
||||
train_idx,
|
||||
val_idx,
|
||||
test_idx,
|
||||
diff_weight,
|
||||
)
|
||||
|
||||
|
||||
def process_dataset_appnp(epsilon):
|
||||
k = 20
|
||||
alpha = 0.2
|
||||
dataset = PubmedGraphDataset()
|
||||
graph = dataset[0]
|
||||
feat = graph.ndata.pop("feat")
|
||||
label = graph.ndata.pop("label")
|
||||
|
||||
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()
|
||||
val_idx = th.nonzero(val_mask, as_tuple=False).squeeze()
|
||||
test_idx = th.nonzero(test_mask, as_tuple=False).squeeze()
|
||||
|
||||
appnp = APPNPConv(k, alpha)
|
||||
id = th.eye(graph.num_nodes()).float()
|
||||
diff_adj = appnp(graph.add_self_loop(), id).numpy()
|
||||
|
||||
diff_adj[diff_adj < epsilon] = 0
|
||||
scaler = MinMaxScaler()
|
||||
scaler.fit(diff_adj)
|
||||
diff_adj = scaler.transform(diff_adj)
|
||||
diff_edges = np.nonzero(diff_adj)
|
||||
diff_weight = diff_adj[diff_edges]
|
||||
diff_graph = dgl.graph(diff_edges)
|
||||
|
||||
return (
|
||||
graph,
|
||||
diff_graph,
|
||||
feat,
|
||||
label,
|
||||
train_idx,
|
||||
val_idx,
|
||||
test_idx,
|
||||
diff_weight,
|
||||
)
|
||||
@@ -0,0 +1,166 @@
|
||||
import argparse
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
from dataset import process_dataset
|
||||
from model import LogReg, MVGRL
|
||||
|
||||
parser = argparse.ArgumentParser(description="mvgrl")
|
||||
|
||||
parser.add_argument(
|
||||
"--dataname", type=str, default="cora", help="Name of dataset."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gpu", type=int, default=0, help="GPU index. Default: -1, using cpu."
|
||||
)
|
||||
parser.add_argument("--epochs", type=int, default=500, help="Training epochs.")
|
||||
parser.add_argument(
|
||||
"--patience",
|
||||
type=int,
|
||||
default=20,
|
||||
help="Patient epochs to wait before early stopping.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lr1", type=float, default=0.001, help="Learning rate of mvgrl."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lr2", type=float, default=0.01, help="Learning rate of linear evaluator."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--wd1", type=float, default=0.0, help="Weight decay of mvgrl."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--wd2", type=float, default=0.0, help="Weight decay of linear evaluator."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--epsilon",
|
||||
type=float,
|
||||
default=0.01,
|
||||
help="Edge mask threshold of diffusion graph.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hid_dim", type=int, default=512, help="Hidden layer dim."
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# check cuda
|
||||
if args.gpu != -1 and th.cuda.is_available():
|
||||
args.device = "cuda:{}".format(args.gpu)
|
||||
else:
|
||||
args.device = "cpu"
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(args)
|
||||
|
||||
# Step 1: Prepare data =================================================================== #
|
||||
(
|
||||
graph,
|
||||
diff_graph,
|
||||
feat,
|
||||
label,
|
||||
train_idx,
|
||||
val_idx,
|
||||
test_idx,
|
||||
edge_weight,
|
||||
) = process_dataset(args.dataname, args.epsilon)
|
||||
n_feat = feat.shape[1]
|
||||
n_classes = np.unique(label).shape[0]
|
||||
|
||||
graph = graph.to(args.device)
|
||||
diff_graph = diff_graph.to(args.device)
|
||||
feat = feat.to(args.device)
|
||||
edge_weight = th.tensor(edge_weight).float().to(args.device)
|
||||
|
||||
train_idx = train_idx.to(args.device)
|
||||
val_idx = val_idx.to(args.device)
|
||||
test_idx = test_idx.to(args.device)
|
||||
|
||||
n_node = graph.num_nodes()
|
||||
lbl1 = th.ones(n_node * 2)
|
||||
lbl2 = th.zeros(n_node * 2)
|
||||
lbl = th.cat((lbl1, lbl2))
|
||||
|
||||
# Step 2: Create model =================================================================== #
|
||||
model = MVGRL(n_feat, args.hid_dim)
|
||||
model = model.to(args.device)
|
||||
|
||||
lbl = lbl.to(args.device)
|
||||
|
||||
# Step 3: Create training components ===================================================== #
|
||||
optimizer = th.optim.Adam(
|
||||
model.parameters(), lr=args.lr1, weight_decay=args.wd1
|
||||
)
|
||||
loss_fn = nn.BCEWithLogitsLoss()
|
||||
|
||||
# Step 4: Training epochs ================================================================ #
|
||||
best = float("inf")
|
||||
cnt_wait = 0
|
||||
for epoch in range(args.epochs):
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
|
||||
shuf_idx = np.random.permutation(n_node)
|
||||
shuf_feat = feat[shuf_idx, :]
|
||||
shuf_feat = shuf_feat.to(args.device)
|
||||
|
||||
out = model(graph, diff_graph, feat, shuf_feat, edge_weight)
|
||||
loss = loss_fn(out, lbl)
|
||||
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
print("Epoch: {0}, Loss: {1:0.4f}".format(epoch, loss.item()))
|
||||
|
||||
if loss < best:
|
||||
best = loss
|
||||
cnt_wait = 0
|
||||
th.save(model.state_dict(), "model.pkl")
|
||||
else:
|
||||
cnt_wait += 1
|
||||
|
||||
if cnt_wait == args.patience:
|
||||
print("Early stopping")
|
||||
break
|
||||
|
||||
model.load_state_dict(th.load("model.pkl"))
|
||||
embeds = model.get_embedding(graph, diff_graph, feat, edge_weight)
|
||||
|
||||
train_embs = embeds[train_idx]
|
||||
test_embs = embeds[test_idx]
|
||||
|
||||
label = label.to(args.device)
|
||||
train_labels = label[train_idx]
|
||||
test_labels = label[test_idx]
|
||||
accs = []
|
||||
|
||||
# Step 5: Linear evaluation ========================================================== #
|
||||
for _ in range(5):
|
||||
model = LogReg(args.hid_dim, n_classes)
|
||||
opt = th.optim.Adam(
|
||||
model.parameters(), lr=args.lr2, weight_decay=args.wd2
|
||||
)
|
||||
|
||||
model = model.to(args.device)
|
||||
loss_fn = nn.CrossEntropyLoss()
|
||||
for epoch in range(300):
|
||||
model.train()
|
||||
opt.zero_grad()
|
||||
logits = model(train_embs)
|
||||
loss = loss_fn(logits, train_labels)
|
||||
loss.backward()
|
||||
opt.step()
|
||||
|
||||
model.eval()
|
||||
logits = model(test_embs)
|
||||
preds = th.argmax(logits, dim=1)
|
||||
acc = th.sum(preds == test_labels).float() / test_labels.shape[0]
|
||||
accs.append(acc * 100)
|
||||
|
||||
accs = th.stack(accs)
|
||||
print(accs.mean().item(), accs.std().item())
|
||||
@@ -0,0 +1,207 @@
|
||||
import argparse
|
||||
import random
|
||||
import warnings
|
||||
|
||||
import dgl
|
||||
|
||||
import numpy as np
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
from dataset import process_dataset, process_dataset_appnp
|
||||
from model import LogReg, MVGRL
|
||||
|
||||
parser = argparse.ArgumentParser(description="mvgrl")
|
||||
|
||||
parser.add_argument(
|
||||
"--dataname", type=str, default="cora", help="Name of dataset."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gpu", type=int, default=-1, help="GPU index. Default: -1, using cpu."
|
||||
)
|
||||
parser.add_argument("--epochs", type=int, default=500, help="Training epochs.")
|
||||
parser.add_argument(
|
||||
"--patience",
|
||||
type=int,
|
||||
default=20,
|
||||
help="Patient epochs to wait before early stopping.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lr1", type=float, default=0.001, help="Learning rate of mvgrl."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lr2", type=float, default=0.01, help="Learning rate of linear evaluator."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--wd1", type=float, default=0.0, help="Weight decay of mvgrl."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--wd2", type=float, default=0.0, help="Weight decay of linear evaluator."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--epsilon",
|
||||
type=float,
|
||||
default=0.01,
|
||||
help="Edge mask threshold of diffusion graph.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hid_dim", type=int, default=512, help="Hidden layer dim."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample_size", type=int, default=2000, help="Subgraph size."
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# check cuda
|
||||
if args.gpu != -1 and th.cuda.is_available():
|
||||
args.device = "cuda:{}".format(args.gpu)
|
||||
else:
|
||||
args.device = "cpu"
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(args)
|
||||
|
||||
# Step 1: Prepare data =================================================================== #
|
||||
if args.dataname == "pubmed":
|
||||
(
|
||||
graph,
|
||||
diff_graph,
|
||||
feat,
|
||||
label,
|
||||
train_idx,
|
||||
val_idx,
|
||||
test_idx,
|
||||
edge_weight,
|
||||
) = process_dataset_appnp(args.epsilon)
|
||||
else:
|
||||
(
|
||||
graph,
|
||||
diff_graph,
|
||||
feat,
|
||||
label,
|
||||
train_idx,
|
||||
val_idx,
|
||||
test_idx,
|
||||
edge_weight,
|
||||
) = process_dataset(args.dataname, args.epsilon)
|
||||
edge_weight = th.tensor(edge_weight).float()
|
||||
graph.ndata["feat"] = feat
|
||||
diff_graph.edata["edge_weight"] = edge_weight
|
||||
|
||||
n_feat = feat.shape[1]
|
||||
n_classes = np.unique(label).shape[0]
|
||||
edge_weight = th.tensor(edge_weight).float()
|
||||
|
||||
train_idx = train_idx.to(args.device)
|
||||
val_idx = val_idx.to(args.device)
|
||||
test_idx = test_idx.to(args.device)
|
||||
|
||||
n_node = graph.num_nodes()
|
||||
|
||||
sample_size = args.sample_size
|
||||
|
||||
lbl1 = th.ones(sample_size * 2)
|
||||
lbl2 = th.zeros(sample_size * 2)
|
||||
lbl = th.cat((lbl1, lbl2))
|
||||
lbl = lbl.to(args.device)
|
||||
|
||||
# Step 2: Create model =================================================================== #
|
||||
model = MVGRL(n_feat, args.hid_dim)
|
||||
model = model.to(args.device)
|
||||
|
||||
# Step 3: Create training components ===================================================== #
|
||||
optimizer = th.optim.Adam(
|
||||
model.parameters(), lr=args.lr1, weight_decay=args.wd1
|
||||
)
|
||||
loss_fn = nn.BCEWithLogitsLoss()
|
||||
|
||||
node_list = list(range(n_node))
|
||||
|
||||
# Step 4: Training epochs ================================================================ #
|
||||
best = float("inf")
|
||||
cnt_wait = 0
|
||||
for epoch in range(args.epochs):
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
|
||||
sample_idx = random.sample(node_list, sample_size)
|
||||
|
||||
g = dgl.node_subgraph(graph, sample_idx)
|
||||
dg = dgl.node_subgraph(diff_graph, sample_idx)
|
||||
|
||||
f = g.ndata.pop("feat")
|
||||
ew = dg.edata.pop("edge_weight")
|
||||
|
||||
shuf_idx = np.random.permutation(sample_size)
|
||||
sf = f[shuf_idx, :]
|
||||
|
||||
g = g.to(args.device)
|
||||
dg = dg.to(args.device)
|
||||
f = f.to(args.device)
|
||||
ew = ew.to(args.device)
|
||||
|
||||
sf = sf.to(args.device)
|
||||
|
||||
out = model(g, dg, f, sf, ew)
|
||||
loss = loss_fn(out, lbl)
|
||||
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
print("Epoch: {0}, Loss: {1:0.4f}".format(epoch, loss.item()))
|
||||
|
||||
if loss < best:
|
||||
best = loss
|
||||
cnt_wait = 0
|
||||
th.save(model.state_dict(), "model.pkl")
|
||||
else:
|
||||
cnt_wait += 1
|
||||
|
||||
if cnt_wait == args.patience:
|
||||
print("Early stopping")
|
||||
break
|
||||
|
||||
model.load_state_dict(th.load("model.pkl"))
|
||||
|
||||
graph = graph.to(args.device)
|
||||
diff_graph = diff_graph.to(args.device)
|
||||
feat = feat.to(args.device)
|
||||
edge_weight = edge_weight.to(args.device)
|
||||
embeds = model.get_embedding(graph, diff_graph, feat, edge_weight)
|
||||
|
||||
train_embs = embeds[train_idx]
|
||||
test_embs = embeds[test_idx]
|
||||
|
||||
label = label.to(args.device)
|
||||
train_labels = label[train_idx]
|
||||
test_labels = label[test_idx]
|
||||
accs = []
|
||||
|
||||
# Step 5: Linear evaluation ========================================================== #
|
||||
for _ in range(5):
|
||||
model = LogReg(args.hid_dim, n_classes)
|
||||
opt = th.optim.Adam(
|
||||
model.parameters(), lr=args.lr2, weight_decay=args.wd2
|
||||
)
|
||||
|
||||
model = model.to(args.device)
|
||||
loss_fn = nn.CrossEntropyLoss()
|
||||
for epoch in range(300):
|
||||
model.train()
|
||||
opt.zero_grad()
|
||||
logits = model(train_embs)
|
||||
loss = loss_fn(logits, train_labels)
|
||||
loss.backward()
|
||||
opt.step()
|
||||
|
||||
model.eval()
|
||||
logits = model(test_embs)
|
||||
preds = th.argmax(logits, dim=1)
|
||||
acc = th.sum(preds == test_labels).float() / test_labels.shape[0]
|
||||
accs.append(acc * 100)
|
||||
|
||||
accs = th.stack(accs)
|
||||
print(accs.mean().item(), accs.std().item())
|
||||
@@ -0,0 +1,74 @@
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
|
||||
from dgl.nn.pytorch import GraphConv
|
||||
from dgl.nn.pytorch.glob import AvgPooling
|
||||
|
||||
|
||||
class LogReg(nn.Module):
|
||||
def __init__(self, hid_dim, n_classes):
|
||||
super(LogReg, self).__init__()
|
||||
|
||||
self.fc = nn.Linear(hid_dim, n_classes)
|
||||
|
||||
def forward(self, x):
|
||||
ret = self.fc(x)
|
||||
return ret
|
||||
|
||||
|
||||
class Discriminator(nn.Module):
|
||||
def __init__(self, dim):
|
||||
super(Discriminator, self).__init__()
|
||||
self.fn = nn.Bilinear(dim, dim, 1)
|
||||
|
||||
def forward(self, h1, h2, h3, h4, c1, c2):
|
||||
c_x1 = c1.expand_as(h1).contiguous()
|
||||
c_x2 = c2.expand_as(h2).contiguous()
|
||||
|
||||
# positive
|
||||
sc_1 = self.fn(h2, c_x1).squeeze(1)
|
||||
sc_2 = self.fn(h1, c_x2).squeeze(1)
|
||||
|
||||
# negative
|
||||
sc_3 = self.fn(h4, c_x1).squeeze(1)
|
||||
sc_4 = self.fn(h3, c_x2).squeeze(1)
|
||||
|
||||
logits = th.cat((sc_1, sc_2, sc_3, sc_4))
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
class MVGRL(nn.Module):
|
||||
def __init__(self, in_dim, out_dim):
|
||||
super(MVGRL, self).__init__()
|
||||
|
||||
self.encoder1 = GraphConv(
|
||||
in_dim, out_dim, norm="both", bias=True, activation=nn.PReLU()
|
||||
)
|
||||
self.encoder2 = GraphConv(
|
||||
in_dim, out_dim, norm="none", bias=True, activation=nn.PReLU()
|
||||
)
|
||||
self.pooling = AvgPooling()
|
||||
|
||||
self.disc = Discriminator(out_dim)
|
||||
self.act_fn = nn.Sigmoid()
|
||||
|
||||
def get_embedding(self, graph, diff_graph, feat, edge_weight):
|
||||
h1 = self.encoder1(graph, feat)
|
||||
h2 = self.encoder2(diff_graph, feat, edge_weight=edge_weight)
|
||||
|
||||
return (h1 + h2).detach()
|
||||
|
||||
def forward(self, graph, diff_graph, feat, shuf_feat, edge_weight):
|
||||
h1 = self.encoder1(graph, feat)
|
||||
h2 = self.encoder2(diff_graph, feat, edge_weight=edge_weight)
|
||||
|
||||
h3 = self.encoder1(graph, shuf_feat)
|
||||
h4 = self.encoder2(diff_graph, shuf_feat, edge_weight=edge_weight)
|
||||
|
||||
c1 = self.act_fn(self.pooling(graph, h1))
|
||||
c2 = self.act_fn(self.pooling(graph, h2))
|
||||
|
||||
out = self.disc(h1, h2, h3, h4, c1, c2)
|
||||
|
||||
return out
|
||||
Reference in New Issue
Block a user