Skip to content

concat

Concatenate several datasets into one flat index space for multi-dataset joint training.

Classes:

ConcatDataset [source] 

ConcatDataset(datasets: Sequence[Dataset])

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 
from torch_pointcloud.datasets import ConcatDataset, S3DIS, ScanNet20

dataset = ConcatDataset([ScanNet20(root="data", split="train"), S3DIS(root="data", areas=["Area_1"])])
len(dataset)  # len(scannet) + len(s3dis)
dataset.sizes  # [len(scannet), len(s3dis)]

SingleDatasetBatchSampler [source] 

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 ConcatDataset order.

  • 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.Generator for shuffling.

Shape

Each yielded value is a List[int] of length batch_size (or fewer for a trailing batch when drop_last is False).

Example 
import torch
from torch.utils.data import DataLoader, TensorDataset

from torch_pointcloud.datasets import ConcatDataset, SingleDatasetBatchSampler

dataset = ConcatDataset([TensorDataset(torch.zeros(16, 3)), TensorDataset(torch.zeros(8, 3))])
sampler = SingleDatasetBatchSampler(dataset.sizes, ratios=[2, 1], batch_size=4)
loader = DataLoader(dataset, batch_sampler=sampler)
len(loader)  # 6: two batches of the first dataset for one of the second, until the first is exhausted