Dataset to concatenate multiple datasets. Purpose: useful to assemble different existing datasets, possibly large-scale datasets as the concatenation operation is done in an on-the-fly manner. Arguments: datasets (sequence): List of datasets to be concatenated.
| 304 | return rtn_ds |
| 305 | |
| 306 | class ConcatDataset(data.Dataset): |
| 307 | """ |
| 308 | Dataset to concatenate multiple datasets. |
| 309 | Purpose: useful to assemble different existing datasets, possibly |
| 310 | large-scale datasets as the concatenation operation is done in an |
| 311 | on-the-fly manner. |
| 312 | Arguments: |
| 313 | datasets (sequence): List of datasets to be concatenated. |
| 314 | """ |
| 315 | |
| 316 | @staticmethod |
| 317 | def cumsum(sequence, weights): |
| 318 | r, s = [], 0 |
| 319 | for i, e in enumerate(sequence): |
| 320 | l = int(len(e) * weights[i]) |
| 321 | r.append(l + s) |
| 322 | s += l |
| 323 | return r |
| 324 | |
| 325 | def __init__(self, datasets, weights=None, **kwargs): |
| 326 | super(ConcatDataset, self).__init__() |
| 327 | assert len(datasets) > 0, 'datasets should not be an empty iterable' |
| 328 | self.datasets = list(datasets) |
| 329 | if weights is None: |
| 330 | self.weights = [1] * len(self.datasets) |
| 331 | else: |
| 332 | self.weights = weights |
| 333 | |
| 334 | assert len(self.weights) == len(self.datasets), 'Length of weights and datasets should be the same.' |
| 335 | self.cumulative_sizes = self.cumsum(self.datasets, self.weights) |
| 336 | |
| 337 | def __len__(self): |
| 338 | return self.cumulative_sizes[-1] |
| 339 | |
| 340 | def __getitem__(self, idx): |
| 341 | dataset_idx = bisect_right(self.cumulative_sizes, idx) |
| 342 | if dataset_idx == 0: |
| 343 | sample_idx = idx |
| 344 | else: |
| 345 | sample_idx = idx - self.cumulative_sizes[dataset_idx - 1] |
| 346 | sample_idx = sample_idx % len(self.datasets[dataset_idx]) |
| 347 | return self.datasets[dataset_idx][sample_idx] |
| 348 | |
| 349 | class RandomMappingDataset(data.Dataset): |
| 350 | ''' |