Fold a batch dim from [bsz, xxx] into [new_bsz, bsz // new_bsz, xxx]
(data: 'DataProto', new_batch_size)
| 110 | |
| 111 | |
| 112 | def fold_batch_dim(data: 'DataProto', new_batch_size): |
| 113 | """ |
| 114 | Fold a batch dim from [bsz, xxx] into [new_bsz, bsz // new_bsz, xxx] |
| 115 | """ |
| 116 | batch_size = data.batch.batch_size[0] |
| 117 | |
| 118 | assert batch_size % new_batch_size == 0 |
| 119 | |
| 120 | tensor: TensorDict = data.batch |
| 121 | non_tensor = data.non_tensor_batch |
| 122 | |
| 123 | tensor = tensor.view(new_batch_size, -1) |
| 124 | tensor.auto_batch_size_(batch_dims=1) |
| 125 | |
| 126 | for key, val in non_tensor.items(): |
| 127 | non_tensor[key] = np.reshape(val, newshape=(new_batch_size, -1, *val.shape[1:])) |
| 128 | |
| 129 | return DataProto(batch=tensor, non_tensor_batch=non_tensor, meta_info=data.meta_info) |
| 130 | |
| 131 | |
| 132 | def unfold_batch_dim(data: 'DataProto', batch_dims=2): |