| 32 | |
| 33 | |
| 34 | class Iterable: |
| 35 | def __init__(self, batch_size=4, shape=(10, 10), epoch_size=None, dtype=None): |
| 36 | self.count = 0 |
| 37 | self.batch_size = batch_size |
| 38 | self.shape = shape |
| 39 | self.epoch_size = epoch_size |
| 40 | self.dtype = dtype or np.int16 |
| 41 | |
| 42 | def __iter__(self): |
| 43 | self.count = 0 |
| 44 | return self |
| 45 | |
| 46 | def __next__(self): |
| 47 | if self.epoch_size is not None and self.epoch_size <= self.count: |
| 48 | raise StopIteration |
| 49 | batch = [ |
| 50 | np.full(self.shape, self.count + i, dtype=self.dtype) for i in range(self.batch_size) |
| 51 | ] |
| 52 | self.count += 1 |
| 53 | return batch |
| 54 | |
| 55 | |
| 56 | class FaultyResetIterable(Iterable): |
no outgoing calls