blob: 7c91caa906f0f62a94cb376ee1dc9418c6f05eae (
plain)
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
|
class DatasetIterator:
"""# Dataset Iterator
Creates an iterator object on a dataset and labels"""
def __init__(self, dataset, labels=None, batch_size=None):
"""# Initialize the iterator with the dataset and labels
- batch_size: How many to include in the iteration"""
self.dataset = dataset
self.labels = labels
self.current = 0
self.batch_size = (
batch_size if batch_size is not None else self.dataset.shape[0]
)
def __iter__(self):
"""# Iterator Function"""
return self
def __next__(self):
"""# Next Iteration
Slice the dataset and labels to provide"""
self.cur = self.current
self.current += 1
if self.cur * self.batch_size < self.dataset.shape[0]:
iterator_slice = slice(
self.cur * self.batch_size, (self.cur + 1) * self.batch_size
)
if self.labels is None:
return self.dataset[iterator_slice, ...]
return self.dataset[iterator_slice, ...], self.labels[iterator_slice, ...]
raise StopIteration
|
