diffengine.datasets.samplers.batch_sampler

Module Contents

Classes

AspectRatioBatchSampler

Aspect ratio batch sampler.

class diffengine.datasets.samplers.batch_sampler.AspectRatioBatchSampler(sampler, batch_size, bucket_ids=None, *, drop_last=False)

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__()

Get the iterator of the sampler.

Return type:

collections.abc.Generator

__len__()

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) –