|
import logging |
|
|
|
import pytest |
|
|
|
from s3prl.dataio.sampler import ( |
|
DistributedBatchSamplerWrapper, |
|
FixedBatchSizeBatchSampler, |
|
MaxTimestampBatchSampler, |
|
) |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
def _merge_batch_indices(batch_indices): |
|
all_indices = [] |
|
for indices in batch_indices: |
|
all_indices += indices |
|
return all_indices |
|
|
|
|
|
@pytest.mark.parametrize("world_size", [1, 2, 3, 4, 5, 6, 7, 8]) |
|
def test_distributed_sampler(world_size): |
|
sampler = [[1, 2, 3], [4, 5, 6, 7], [8], [9, 10]] |
|
ddp_indices = [] |
|
for rank in range(world_size): |
|
ddp_sampler = DistributedBatchSamplerWrapper(sampler, world_size, rank) |
|
ddp_indices += _merge_batch_indices(ddp_sampler) |
|
assert sorted(ddp_indices) == sorted(_merge_batch_indices(sampler)) |
|
|
|
|
|
timestamps = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] |
|
data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] |
|
|
|
|
|
@pytest.mark.parametrize("batch_size", [1, 2, 3, len(data)]) |
|
def test_FixedBatchSizeBatchSampler(batch_size): |
|
dataset = data |
|
iter1 = list(iter(FixedBatchSizeBatchSampler(dataset, batch_size, shuffle=False))) |
|
iter2 = list(iter(FixedBatchSizeBatchSampler(dataset, batch_size, shuffle=True))) |
|
indices1 = sorted(_merge_batch_indices(iter1)) |
|
indices2 = sorted(_merge_batch_indices(iter2)) |
|
assert indices1 == indices2 == list(range(len(timestamps))) |
|
|
