def generator(self):
return batch_generator(
self.dataset,
batch_size=self.batch_size,
epochs=self.epochs,
shuffle=self.shuffle,
)
def generator(self):
"""
Returns lists (batches) of `Graph` objects.
"""
return batch_generator(
self.dataset,
batch_size=self.batch_size,
epochs=self.epochs,
shuffle=self.shuffle,
)
def test_batch_generator():
size = 10
batch_size = 6
a = list(range(size))
b = np.arange(size)
class TestDataset(Dataset):
def read(self):
return [
Graph(x=np.random.rand(n, 2),
a=np.random.randint(0, 2, (n, n)),
y=np.array([0., 1.])) for n in range(size)
]
c = TestDataset()
batches = batch_generator([a, b, c], batch_size=batch_size, epochs=10)
for batch in batches:
a_, b_, c_ = batch
for i in range(len(a_)):
assert a_[i] == b_[i] == c_[i].n_nodes
