(
self,
dataset,
batch_size=1,
drop_last=False,
num_samples=None,
world_size=None,
rank=None,
seed=None,
)
| 35 | """ |
| 36 | |
| 37 | def __init__( |
| 38 | self, |
| 39 | dataset, |
| 40 | batch_size=1, |
| 41 | drop_last=False, |
| 42 | num_samples=None, |
| 43 | world_size=None, |
| 44 | rank=None, |
| 45 | seed=None, |
| 46 | ): |
| 47 | if ( |
| 48 | not isinstance(batch_size, int) |
| 49 | or isinstance(batch_size, bool) |
| 50 | or batch_size <= 0 |
| 51 | ): |
| 52 | raise ValueError( |
| 53 | "batch_size should be a positive integer value, " |
| 54 | "but got batch_size={}".format(batch_size) |
| 55 | ) |
| 56 | if not isinstance(drop_last, bool): |
| 57 | raise ValueError( |
| 58 | "drop_last should be a boolean value, but got " |
| 59 | "drop_last={}".format(drop_last) |
| 60 | ) |
| 61 | |
| 62 | if num_samples is not None and ( |
| 63 | not isinstance(num_samples, int) or num_samples <= 0 |
| 64 | ): |
| 65 | raise ValueError( |
| 66 | "num_samples should be a positive integer " |
| 67 | "value, but got num_samples={}".format(num_samples) |
| 68 | ) |
| 69 | |
| 70 | self.batch_size = batch_size |
| 71 | self.dataset = dataset |
| 72 | self.drop_last = drop_last |
| 73 | |
| 74 | if world_size is None: |
| 75 | world_size = dist.get_world_size() if dist.is_distributed() else 1 |
| 76 | self.world_size = world_size |
| 77 | if rank is None: |
| 78 | rank = dist.get_rank() if dist.is_distributed() else 0 |
| 79 | self.rank = rank |
| 80 | |
| 81 | if num_samples is None: |
| 82 | num_samples = len(self.dataset) |
| 83 | self.num_samples = int(math.ceil(num_samples / self.world_size)) |
| 84 | |
| 85 | if self.num_samples < self.batch_size: |
| 86 | raise ValueError( |
| 87 | "num_samples should be greater than batch_size " |
| 88 | ", but got num_samples={} and batch_size={}".format( |
| 89 | self.num_samples, self.batch_size |
| 90 | ) |
| 91 | ) |
| 92 | |
| 93 | # Make sure seeds are the same at each rank |
| 94 | if seed is None and self.world_size > 1: |
nothing calls this directly
no test coverage detected