concat
Concatenate several datasets into one flat index space for multi-dataset joint training.
Classes:
-
ConcatDataset–Concatenates several datasets into one flat index space.
-
SingleDatasetBatchSampler–Yields batches whose global indices all come from one dataset of a
ConcatDataset.
ConcatDataset
¶
Bases: Dataset
Concatenates several datasets into one flat index space.
Each child dataset keeps its own transform, so datasets from different domains (e.g. ScanNet and
S3DIS) train jointly while stamping their own condition key and mapping to their native label space.
Pair it with SingleDatasetBatchSampler to keep every batch single-domain so per-dataset
normalization statistics (BatchNorm, PDNorm) stay clean.
Parameters:
-
datasets(Sequence[Dataset]) –The datasets to concatenate, in order. The first is treated as the main dataset by
SingleDatasetBatchSampler(its exhaustion ends the epoch).
Example
SingleDatasetBatchSampler
¶
SingleDatasetBatchSampler(
sizes: Sequence[int],
ratios: Sequence[int],
batch_size: int,
shuffle: bool = True,
drop_last: bool = True,
generator: Optional[Generator] = None,
)
Bases: Sampler[List[int]]
Yields batches whose global indices all come from one dataset of a ConcatDataset.
Given the per-dataset sizes and one positive integer ratio per dataset, it partitions each dataset's
indices into batches of batch_size and interleaves the datasets round-robin weighted by the ratios.
Within a round the first dataset yields ratios[0] batches, the second ratios[1], and so on. The
first (main) dataset drives the epoch length: when it
is exhausted the epoch ends, while the other datasets restart (reshuffled) as needed. Because every
yielded batch is drawn from a single dataset, per-batch normalization (BatchNorm, PDNorm) sees a
single domain.
Parameters:
-
sizes(Sequence[int]) –Number of samples in each child dataset, in
ConcatDatasetorder. -
ratios(Sequence[int]) –One positive integer sampling weight per dataset, aligned with
sizes. -
batch_size(int) –Number of indices per batch.
-
shuffle(bool, default:True) –Shuffle each dataset's indices, reshuffling on restart.
-
drop_last(bool, default:True) –Drop each dataset's trailing partial batch.
-
generator(Optional[Generator], default:None) –Optional
torch.Generatorfor shuffling.
Shape
Each yielded value is a List[int] of length batch_size (or fewer for a trailing batch when
drop_last is False).