|
from collections import Counter |
|
|
|
import pytest |
|
|
|
from s3prl.dataio.sampler import BalancedWeightedSampler |
|
|
|
|
|
@pytest.mark.parametrize("duplicate", [10000, 100000]) |
|
def test_balanced_weighted_sampler(duplicate: int): |
|
labels = ["a", "a", "b", "a"] |
|
batch_size = 5 |
|
prev_diff_ratio = 1.0 |
|
sampler = BalancedWeightedSampler( |
|
labels, batch_size=batch_size, duplicate=duplicate, seed=0 |
|
) |
|
indices = list(sampler) |
|
assert len(indices[0]) == batch_size |
|
|
|
counter = Counter() |
|
for batch_indices in indices: |
|
for idx in batch_indices: |
|
counter.update(labels[idx]) |
|
|
|
diff_ratio = abs(counter["a"] - counter["b"]) / duplicate * len(labels) |
|
assert diff_ratio < prev_diff_ratio |
|
prev_diff_ratio = diff_ratio |
|
|
|
diff_ratio < 0.05 |
|
|
