(batch_size: int, *args)
| 96 | |
| 97 | |
| 98 | def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]: |
| 99 | assert len(args) > 0 and all( |
| 100 | len(a) == len(args[0]) for a in args |
| 101 | ), "Batched iteration must have inputs of all the same size." |
| 102 | n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0) |
| 103 | for b in range(n_batches): |
| 104 | yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args] |
| 105 | |
| 106 | |
| 107 | def mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]: |