| 148 | |
| 149 | |
| 150 | class _ExternalSourceGroup(object): |
| 151 | def __init__( |
| 152 | self, |
| 153 | callback, |
| 154 | source_desc, |
| 155 | is_multioutput, |
| 156 | instances=[], |
| 157 | *, |
| 158 | cuda_stream=None, |
| 159 | use_copy_kernel=None, |
| 160 | batch=True, |
| 161 | parallel=False, |
| 162 | prefetch_queue_depth=None, |
| 163 | bytes_per_sample_hint=None, |
| 164 | batch_info=None, |
| 165 | ): |
| 166 | self.instances = list(instances) # we need a copy! |
| 167 | self.utilized_instances = self.instances |
| 168 | self.is_multioutput = is_multioutput |
| 169 | self.callback = callback |
| 170 | self.source_desc = source_desc |
| 171 | self._cuda_stream = cuda_stream |
| 172 | self.use_copy_kernel = use_copy_kernel |
| 173 | self.batch = batch |
| 174 | self.batch_info = batch_info |
| 175 | # Index of a batch within the epoch that will be returned from |
| 176 | # get_batch or schedule_and_receive call. Contrary to Pipeline's |
| 177 | # `epoch_idx` it is tracked separately by ExternalSourceGroup due to |
| 178 | # prefetching of batches in parallel mode. |
| 179 | self.current_iter = 0 |
| 180 | self.current_sample = 0 |
| 181 | self.parallel = parallel |
| 182 | self.prefetch_queue_depth = prefetch_queue_depth |
| 183 | self.bytes_per_sample_hint = bytes_per_sample_hint |
| 184 | if callback is not None: |
| 185 | self.accepts_arg = _accepted_arg_count(callback) > 0 |
| 186 | |
| 187 | def append(self, instance): |
| 188 | self.instances.append(instance) |
| 189 | self.utilized_instances = self.instances |
| 190 | |
| 191 | def feed_count(self, pipeline): |
| 192 | return pipeline.input_feed_count(self.utilized_instances[0]._name) |
| 193 | |
| 194 | def disable_pruned_instances(self, pruned_mask): |
| 195 | if len(pruned_mask) != len(self.instances): |
| 196 | raise RuntimeError( |
| 197 | f"Mask of the pruned outputs of the external source must have the length matching " |
| 198 | f"the number of outputs of the external source. The external source node has " |
| 199 | f"{len(self.instances)} outputs, but received mask of length {len(pruned_mask)}." |
| 200 | ) |
| 201 | self.utilized_instances = [ |
| 202 | instance for instance, is_pruned in zip(self.instances, pruned_mask) if not is_pruned |
| 203 | ] |
| 204 | |
| 205 | def callback_args(self, idx_in_batch, epoch_idx, batch_size=0, lead=0): |
| 206 | """Generate information to be passed to ES callback. |
| 207 | |