| 194 | batch_size = 7 |
| 195 | |
| 196 | class IterSetupPipeline(Pipeline): |
| 197 | def __init__(self, iterator, num_threads, device_id, device): |
| 198 | super(IterSetupPipeline, self).__init__( |
| 199 | batch_size=iterator.batch_size, num_threads=num_threads, device_id=device_id |
| 200 | ) |
| 201 | |
| 202 | self.iterator = iterator |
| 203 | self._device = device |
| 204 | |
| 205 | def define_graph(self): |
| 206 | if use_fn_api: |
| 207 | self.batch_1 = fn.external_source(device=self._device, name="src1") |
| 208 | self.batch_2 = fn.external_source(device=self._device, name="src2") |
| 209 | else: |
| 210 | input_1 = ops.ExternalSource(device=self._device) |
| 211 | input_2 = ops.ExternalSource(device=self._device) |
| 212 | self.batch_1 = input_1(name="src1") |
| 213 | self.batch_2 = input_2(name="src2") |
| 214 | return [self.batch_1, self.batch_2] |
| 215 | |
| 216 | def iter_setup(self): |
| 217 | batch_1, batch_2 = next(self.iterator) |
| 218 | if by_name: |
| 219 | self.feed_input("src1", batch_1) |
| 220 | self.feed_input("src2", batch_2) |
| 221 | else: |
| 222 | self.feed_input(self.batch_1, batch_1) |
| 223 | self.feed_input(self.batch_2, batch_2) |
| 224 | |
| 225 | iter_num = 5 |
| 226 | source = TestIterator(iter_num, batch_size, [2, 3]) |
no outgoing calls