| 47 | return optimizer |
| 48 | |
| 49 | def build_dataloader(self, dataset, shuffle, max_tokens=None, max_sentences=None, |
| 50 | required_batch_size_multiple=-1, endless=False, batch_by_size=True): |
| 51 | devices_cnt = torch.cuda.device_count() |
| 52 | if devices_cnt == 0: |
| 53 | devices_cnt = 1 |
| 54 | if required_batch_size_multiple == -1: |
| 55 | required_batch_size_multiple = devices_cnt |
| 56 | |
| 57 | def shuffle_batches(batches): |
| 58 | np.random.shuffle(batches) |
| 59 | return batches |
| 60 | |
| 61 | if max_tokens is not None: |
| 62 | max_tokens *= devices_cnt |
| 63 | if max_sentences is not None: |
| 64 | max_sentences *= devices_cnt |
| 65 | indices = dataset.ordered_indices() |
| 66 | if batch_by_size: |
| 67 | batch_sampler = utils.batch_by_size( |
| 68 | indices, dataset.num_tokens, max_tokens=max_tokens, max_sentences=max_sentences, |
| 69 | required_batch_size_multiple=required_batch_size_multiple, |
| 70 | ) |
| 71 | else: |
| 72 | batch_sampler = [] |
| 73 | for i in range(0, len(indices), max_sentences): |
| 74 | batch_sampler.append(indices[i:i + max_sentences]) |
| 75 | |
| 76 | if shuffle: |
| 77 | batches = shuffle_batches(list(batch_sampler)) |
| 78 | if endless: |
| 79 | batches = [b for _ in range(1000) for b in shuffle_batches(list(batch_sampler))] |
| 80 | else: |
| 81 | batches = batch_sampler |
| 82 | if endless: |
| 83 | batches = [b for _ in range(1000) for b in batches] |
| 84 | num_workers = dataset.num_workers |
| 85 | if self.trainer.use_ddp: |
| 86 | num_replicas = dist.get_world_size() |
| 87 | rank = dist.get_rank() |
| 88 | batches = [x[rank::num_replicas] for x in batches if len(x) % num_replicas == 0] |
| 89 | return torch.utils.data.DataLoader(dataset, |
| 90 | collate_fn=dataset.collater, |
| 91 | batch_sampler=batches, |
| 92 | num_workers=num_workers, |
| 93 | pin_memory=False) |
| 94 | |
| 95 | def build_phone_encoder(self, data_dir): |
| 96 | phone_list_file = os.path.join(data_dir, 'phone_set.json') |