(self, input: List[Input], batch_size,
**kwargs)
| 302 | return batch_data |
| 303 | |
| 304 | def _process_batch(self, input: List[Input], batch_size, |
| 305 | **kwargs) -> Dict[str, Any]: |
| 306 | preprocess_params = kwargs.get('preprocess_params') |
| 307 | forward_params = kwargs.get('forward_params') |
| 308 | postprocess_params = kwargs.get('postprocess_params') |
| 309 | |
| 310 | # batch data |
| 311 | output_list = [] |
| 312 | for i in range(0, len(input), batch_size): |
| 313 | end = min(i + batch_size, len(input)) |
| 314 | real_batch_size = end - i |
| 315 | preprocessed_list = [ |
| 316 | self.preprocess(i, **preprocess_params) for i in input[i:end] |
| 317 | ] |
| 318 | |
| 319 | with device_placement(self.framework, self.device_name): |
| 320 | if self.framework == Frameworks.torch: |
| 321 | with torch.no_grad(): |
| 322 | batched_out = self._batch(preprocessed_list) |
| 323 | if self._auto_collate: |
| 324 | batched_out = self._collate_fn(batched_out) |
| 325 | batched_out = self.forward(batched_out, |
| 326 | **forward_params) |
| 327 | else: |
| 328 | batched_out = self._batch(preprocessed_list) |
| 329 | batched_out = self.forward(batched_out, **forward_params) |
| 330 | |
| 331 | for batch_idx in range(real_batch_size): |
| 332 | out = {} |
| 333 | for k, element in batched_out.items(): |
| 334 | if element is not None: |
| 335 | if isinstance(element, (tuple, list)): |
| 336 | if isinstance(element[0], torch.Tensor): |
| 337 | out[k] = type(element)( |
| 338 | e[batch_idx:batch_idx + 1] |
| 339 | for e in element) |
| 340 | else: |
| 341 | # Compatible with traditional pipelines |
| 342 | out[k] = element[batch_idx] |
| 343 | else: |
| 344 | out[k] = element[batch_idx:batch_idx + 1] |
| 345 | out = self.postprocess(out, **postprocess_params) |
| 346 | self._check_output(out) |
| 347 | output_list.append(out) |
| 348 | |
| 349 | return output_list |
| 350 | |
| 351 | def _check_input(self, input): |
| 352 | task_name = self.group_key |
no test coverage detected