| 151 | |
| 152 | |
| 153 | class SampleIterator: |
| 154 | def __init__(self, batch_iterator, is_multioutput=False): |
| 155 | self.src = batch_iterator |
| 156 | self.is_multioutput = is_multioutput |
| 157 | self.batch = ([],) if is_multioutput else [] |
| 158 | self.idx = 0 |
| 159 | |
| 160 | def __iter__(self): |
| 161 | return SampleIterator(iter(self.src), self.is_multioutput) |
| 162 | |
| 163 | def __next__(self): |
| 164 | batch_size = len(self.batch[0]) if self.is_multioutput else len(self.batch) |
| 165 | if self.idx >= batch_size: |
| 166 | self.idx = 0 |
| 167 | self.batch = next(self.src) |
| 168 | if self.is_multioutput: |
| 169 | ret = tuple(b[self.idx] for b in self.batch) |
| 170 | else: |
| 171 | ret = self.batch[self.idx] |
| 172 | self.idx += 1 |
| 173 | return ret |
| 174 | |
| 175 | next = __next__ |
| 176 | |
| 177 | |
| 178 | def run_and_check(pipe, ref_iterable): |
no outgoing calls