Repeat the batch data a specified number of times. Args: repeat_times (int): Number of times to repeat the data. interleave (bool): Whether to interleave the repeated data. Returns: DataProto: A new DataProto with repeated data.
(self, repeat_times=2, interleave=True)
| 552 | self.non_tensor_batch = {key: val[indices_np] for key, val in self.non_tensor_batch.items()} |
| 553 | |
| 554 | def repeat(self, repeat_times=2, interleave=True): |
| 555 | """ |
| 556 | Repeat the batch data a specified number of times. |
| 557 | |
| 558 | Args: |
| 559 | repeat_times (int): Number of times to repeat the data. |
| 560 | interleave (bool): Whether to interleave the repeated data. |
| 561 | |
| 562 | Returns: |
| 563 | DataProto: A new DataProto with repeated data. |
| 564 | """ |
| 565 | if self.batch is not None: |
| 566 | if interleave: |
| 567 | # Interleave the data |
| 568 | repeated_tensors = { |
| 569 | key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self.batch.items() |
| 570 | } |
| 571 | else: |
| 572 | # Stack the data |
| 573 | repeated_tensors = { |
| 574 | key: tensor.unsqueeze(0).expand(repeat_times, *tensor.shape).reshape(-1, *tensor.shape[1:]) |
| 575 | for key, tensor in self.batch.items() |
| 576 | } |
| 577 | |
| 578 | repeated_batch = TensorDict( |
| 579 | source=repeated_tensors, |
| 580 | batch_size=(self.batch.batch_size[0] * repeat_times,), |
| 581 | ) |
| 582 | else: |
| 583 | repeated_batch = None |
| 584 | |
| 585 | repeated_non_tensor_batch = {} |
| 586 | for key, val in self.non_tensor_batch.items(): |
| 587 | if interleave: |
| 588 | repeated_non_tensor_batch[key] = np.repeat(val, repeat_times, axis=0) |
| 589 | else: |
| 590 | repeated_non_tensor_batch[key] = np.tile(val, (repeat_times,) + (1,) * (val.ndim - 1)) |
| 591 | |
| 592 | return DataProto( |
| 593 | batch=repeated_batch, |
| 594 | non_tensor_batch=repeated_non_tensor_batch, |
| 595 | meta_info=self.meta_info, |
| 596 | ) |
| 597 | |
| 598 | |
| 599 | import ray |