-
Notifications
You must be signed in to change notification settings - Fork 3
/
datasets.py
110 lines (95 loc) · 4.11 KB
/
datasets.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
#!/usr/bin/env python3
from pathlib import Path
import numpy as np
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
from torch_geometric.datasets import TUDataset
from alkane import AlkaneCarbonyl
DATA_DIR = Path(__file__).parent / "data/TUDataset"
# Helper functions ----------
def _append_idx(data):
"Appends an index to each graph to track which explanations are used from the batch"
idx = [
Data(
x=data[i].x,
edge_index=data[i].edge_index,
edge_attr=data[i].edge_attr,
y=data[i].y,
idx=i,
)
for i in range(len(data))
]
return idx
def _split_tu_dataset(dataset, seed, batch_size) -> DataLoader:
num_training = int(len(dataset) * 0.8)
num_val = int(len(dataset) * 0.1)
num_test = len(dataset) - (num_training + num_val)
np.random.seed(seed)
train_dataset = dataset[:num_training]
val_dataset = dataset[num_training : (num_training + num_val)]
test_dataset = dataset[(num_training + num_val) :]
train_data = _append_idx(train_dataset)
val_data = _append_idx(val_dataset)
test_data = _append_idx(test_dataset)
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)
return train_loader, val_loader, test_loader
def _split_molecule_dataset(dataset, seed, batch_size) -> DataLoader:
train_loader = dataset.get_train_loader(batch_size=batch_size)
test_loader = dataset.get_test_loader()
val_loader = dataset.get_val_loader()
train_data = _append_idx(train_loader.dataset)
val_data = _append_idx(val_loader.dataset)
test_data = _append_idx(test_loader.dataset)
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)
return train_loader, val_loader, test_loader
# Data loading functions ----
def load_mutag(seed: int, batch_size: int, split_train_val_test: bool = True):
dataset = TUDataset(root=DATA_DIR, name="MUTAG")
dataset = dataset.shuffle()
if not split_train_val_test:
return DataLoader(dataset)
train_dataset = dataset[:150]
test_dataset = dataset[150:]
np.random.seed(seed)
train_data = _append_idx(train_dataset)
test_data = _append_idx(test_dataset)
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)
val_loader = test_loader
return train_loader, val_loader, test_loader
def load_dd(seed: int, batch_size: int, split_train_val_test: bool = True):
dataset = TUDataset(root=DATA_DIR, name="DD")
dataset = dataset.shuffle()
if split_train_val_test:
train_loader, val_loader, test_loader = _split_tu_dataset(
dataset=dataset, seed=seed, batch_size=batch_size,
)
return train_loader, val_loader, test_loader
return DataLoader(dataset)
def load_protein(seed: int, batch_size: int, split_train_val_test: bool = True):
dataset = TUDataset(root=DATA_DIR, name="PROTEINS_full")
dataset = dataset.shuffle()
if split_train_val_test:
train_loader, val_loader, test_loader = _split_tu_dataset(
dataset=dataset, seed=seed, batch_size=batch_size,
)
return train_loader, val_loader, test_loader
return DataLoader(dataset)
def load_alkane(seed: int, batch_size: int, split_train_val_test: bool = True):
dataset = AlkaneCarbonyl(split_sizes=(0.8, 0.1, 0.1), downsample_seed=seed)
if split_train_val_test:
train_loader, val_loader, test_loader = _split_molecule_dataset(
dataset=dataset, seed=seed, batch_size=batch_size
)
return train_loader, val_loader, test_loader
return DataLoader(dataset)
DATASET_LOADERS = {
"mutag": load_mutag,
"DD": load_dd,
"PROTEIN": load_protein,
"alkane": load_alkane,
}