diffengine.datasets.samplers

Submodules

Package Contents

Classes

AspectRatioBatchSampler

Aspect ratio batch sampler.
class diffengine.datasets.samplers.AspectRatioBatchSampler(sampler, batch_size, bucket_ids=None, *, drop_last=False)[source]

Bases: torch.utils.data.BatchSampler

Aspect ratio batch sampler.

A sampler wrapper for grouping images with similar aspect ratio into a same batch.

Args:

sampler (Sampler): Base sampler. batch_size (int): Size of mini-batch. bucket_ids (str | None): The path of bucket ids. If None, the

bucket ids will be calculated automatically. Default: None.

drop_last (bool): If True, the sampler will drop the last batch if

its size would be less than batch_size.

__iter__()[source]

Get the iterator of the sampler.

Return type:

collections.abc.Generator

__len__()[source]

Get the length of the sampler.

Return type:

int

Parameters:

  • sampler (torch.utils.data.Sampler) –

  • batch_size (int) –

  • bucket_ids (str | None) –

  • drop_last (bool) –