wehub-resource-sync
99 lines
3.2 KiB
Python
99 lines
3.2 KiB
Python
import collections
|
|
|
|
import dgl
|
|
from dgl.data import PPIDataset
|
|
|
|
from torch.utils.data import DataLoader, Dataset
|
|
|
|
# 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 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
|
|
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"]
|
|
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
|
|
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
|
|
|
|
etypes = train_set[0].etypes
|
|
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,
|
|
)
|
|
return train_loader, valid_loader, test_loader, etypes, in_size, out_size
|