wehub-resource-sync
114 lines
4.0 KiB
Python
114 lines
4.0 KiB
Python
import dgl
|
|
import dgl.nn as dglnn
|
|
import sklearn.linear_model as lm
|
|
import sklearn.metrics as skm
|
|
import torch as th
|
|
import torch.functional as F
|
|
import torch.nn as nn
|
|
import tqdm
|
|
|
|
|
|
class SAGE(nn.Module):
|
|
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
|
|
):
|
|
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"))
|
|
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"))
|
|
else:
|
|
self.layers.append(dglnn.SAGEConv(in_feats, n_classes, "mean"))
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.activation = activation
|
|
|
|
def forward(self, blocks, x):
|
|
h = x
|
|
for l, (layer, block) in enumerate(zip(self.layers, blocks)):
|
|
h = layer(block, h)
|
|
if l != len(self.layers) - 1:
|
|
h = self.activation(h)
|
|
h = self.dropout(h)
|
|
return h
|
|
|
|
def inference(self, g, x, device, batch_size, num_workers):
|
|
"""
|
|
Inference with the GraphSAGE model on full neighbors (i.e. without neighbor sampling).
|
|
g : the entire graph.
|
|
x : the input of entire node set.
|
|
|
|
The inference code is written in a fashion that it could handle any number of nodes and
|
|
layers.
|
|
"""
|
|
# During inference with sampling, multi-layer blocks are very inefficient because
|
|
# lots of computations in the first few layers are repeated.
|
|
# Therefore, we compute the representation of all nodes layer by layer. The nodes
|
|
# 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,
|
|
)
|
|
|
|
sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
|
|
dataloader = dgl.dataloading.DataLoader(
|
|
g,
|
|
th.arange(g.num_nodes()).to(g.device),
|
|
sampler,
|
|
device=device if num_workers == 0 else None,
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
drop_last=False,
|
|
num_workers=num_workers,
|
|
)
|
|
|
|
for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
|
|
block = blocks[0]
|
|
|
|
block = block.int().to(device)
|
|
h = x[input_nodes].to(device)
|
|
h = layer(block, h)
|
|
if l != len(self.layers) - 1:
|
|
h = self.activation(h)
|
|
h = self.dropout(h)
|
|
|
|
y[output_nodes] = h.cpu()
|
|
|
|
x = y
|
|
return y
|
|
|
|
|
|
def compute_acc_unsupervised(emb, labels, train_nids, val_nids, test_nids):
|
|
"""
|
|
Compute the accuracy of prediction given the labels.
|
|
"""
|
|
emb = emb.cpu().numpy()
|
|
labels = labels.cpu().numpy()
|
|
train_nids = train_nids.cpu().numpy()
|
|
train_labels = labels[train_nids]
|
|
val_nids = val_nids.cpu().numpy()
|
|
val_labels = labels[val_nids]
|
|
test_nids = test_nids.cpu().numpy()
|
|
test_labels = labels[test_nids]
|
|
|
|
emb = (emb - emb.mean(0, keepdims=True)) / emb.std(0, keepdims=True)
|
|
|
|
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")
|
|
return f1_micro_eval, f1_micro_test
|