Fetch a batch of data without waiting
(self)
| 107 | f.set_result([k[idx] for k in outputs]) |
| 108 | |
| 109 | def fetch_batch(self): |
| 110 | """ Fetch a batch of data without waiting""" |
| 111 | inp, f = self.queue.get() |
| 112 | nr_input_var = len(inp) |
| 113 | batched, futures = [[] for _ in range(nr_input_var)], [] |
| 114 | for k in range(nr_input_var): |
| 115 | batched[k].append(inp[k]) |
| 116 | futures.append(f) |
| 117 | while len(futures) < self.batch_size: |
| 118 | try: |
| 119 | inp, f = self.queue.get_nowait() |
| 120 | for k in range(nr_input_var): |
| 121 | batched[k].append(inp[k]) |
| 122 | futures.append(f) |
| 123 | except queue.Empty: |
| 124 | break # do not wait |
| 125 | |
| 126 | for k in range(nr_input_var): |
| 127 | batched[k] = np.asarray(batched[k]) |
| 128 | return batched, futures |
| 129 | |
| 130 | |
| 131 | class MultiThreadAsyncPredictor(AsyncPredictorBase): |