Test the speed of a DataFlow
| 27 | |
| 28 | |
| 29 | class TestDataSpeed(ProxyDataFlow): |
| 30 | """ Test the speed of a DataFlow """ |
| 31 | def __init__(self, ds, size=5000, warmup=0): |
| 32 | """ |
| 33 | Args: |
| 34 | ds (DataFlow): the DataFlow to test. |
| 35 | size (int): number of datapoints to fetch. |
| 36 | warmup (int): warmup iterations |
| 37 | """ |
| 38 | super(TestDataSpeed, self).__init__(ds) |
| 39 | self.test_size = int(size) |
| 40 | self.warmup = int(warmup) |
| 41 | self._reset_called = False |
| 42 | |
| 43 | def reset_state(self): |
| 44 | self._reset_called = True |
| 45 | super(TestDataSpeed, self).reset_state() |
| 46 | |
| 47 | def __iter__(self): |
| 48 | """ Will run testing at the beginning, then produce data normally. """ |
| 49 | self.start() |
| 50 | yield from self.ds |
| 51 | |
| 52 | def start(self): |
| 53 | """ |
| 54 | Start testing with a progress bar. |
| 55 | """ |
| 56 | if not self._reset_called: |
| 57 | self.ds.reset_state() |
| 58 | itr = self.ds.__iter__() |
| 59 | if self.warmup: |
| 60 | for _ in tqdm.trange(self.warmup, **get_tqdm_kwargs()): |
| 61 | next(itr) |
| 62 | # add smoothing for speed benchmark |
| 63 | with get_tqdm(total=self.test_size, |
| 64 | leave=True, smoothing=0.2) as pbar: |
| 65 | for idx, dp in enumerate(itr): |
| 66 | pbar.update() |
| 67 | if idx == self.test_size - 1: |
| 68 | break |
| 69 | |
| 70 | |
| 71 | class BatchData(ProxyDataFlow): |
no outgoing calls
no test coverage detected