(self, tasks, datasets, reweight=True, temperature=0.8, max_limit=200000)
| 34 | |
| 35 | class MultiTaskDataset(torch.utils.data.Dataset): |
| 36 | def __init__(self, tasks, datasets, reweight=True, temperature=0.8, max_limit=200000): |
| 37 | super(MultiTaskDataset, self).__init__() |
| 38 | self.tasks = tasks |
| 39 | self.datasets = datasets |
| 40 | self.reweight = reweight |
| 41 | self.temperature = temperature |
| 42 | self.lens = [len(dataset) for dataset in datasets] |
| 43 | self.weights = np.array([min(l, max_limit) ** temperature for l in self.lens]) |
| 44 | self.total_len = sum(self.lens) |
| 45 | self.cumulative_lens = list(accumulate(self.lens)) |
| 46 | if self.reweight: |
| 47 | print_rank_0(list(zip(self.tasks, self.lens, self.weights))) |
| 48 | else: |
| 49 | print_rank_0(list(zip(self.tasks, self.lens))) |
| 50 | self.weights /= self.weights.sum() |
| 51 | |
| 52 | def __len__(self): |
| 53 | return self.total_len * 1000 |
nothing calls this directly
no test coverage detected