| 246 | class FinetuneDistSampler(Sampler): |
| 247 | # Distrubuted Sampler ensuring data in a batch are of the same type (e.g. text, image-text) |
| 248 | def __init__(self, dataset: FinetuneDataset, num_replicas: Optional[int] = None, |
| 249 | rank: Optional[int] = None, shuffle: bool = True, |
| 250 | seed: int = 0, batch_size = None, acc_grad=1) -> None: |
| 251 | if num_replicas is None or rank is None or rank >= num_replicas or rank < 0: |
| 252 | raise ValueError( |
| 253 | f"Invalid num_replicas ({num_replicas}) or rank ({rank})") |
| 254 | assert batch_size is not None |
| 255 | self.batch_size = batch_size |
| 256 | |
| 257 | self.dataset = dataset |
| 258 | self.num_replicas = num_replicas |
| 259 | self.rank = rank |
| 260 | self.acc_grad = acc_grad |
| 261 | self.epoch = 0 |
| 262 | self.start_iter = 0 |
| 263 | |
| 264 | group_indices = dataset.groups() |
| 265 | global_bsz = batch_size * num_replicas * acc_grad |
| 266 | len_groups = [len(_) // global_bsz * global_bsz for _ in group_indices] |
| 267 | group_indices = [indices[:len_indices] for indices, len_indices in zip(group_indices, len_groups)] |
| 268 | group_n_batch = [len(_)//batch_size for _ in group_indices] |
| 269 | assert all([_%num_replicas==0 for _ in group_n_batch]) |
| 270 | n_total_batch = sum(group_n_batch) |
| 271 | |
| 272 | assert n_total_batch % self.num_replicas == 0 |
| 273 | |
| 274 | self.group_indices = group_indices |
| 275 | |
| 276 | self.total_size = n_total_batch * batch_size |
| 277 | self.num_samples = self.total_size // num_replicas |
| 278 | self.shuffle = shuffle |
| 279 | self.seed = seed |
| 280 | |
| 281 | def __iter__(self) -> Iterator: |
| 282 | global_batch_size = self.batch_size * self.num_replicas * self.acc_grad |