File size: 1,500 Bytes
07d7c23 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
# Copyright (c) Facebook, Inc. and its affiliates.
import random
import unittest
from typing import Any, Iterable, Iterator, Tuple
from densepose.data import CombinedDataLoader
def _grouper(iterable: Iterable[Any], n: int, fillvalue=None) -> Iterator[Tuple[Any]]:
"""
Group elements of an iterable by chunks of size `n`, e.g.
grouper(range(9), 4) ->
(0, 1, 2, 3), (4, 5, 6, 7), (8, None, None, None)
"""
it = iter(iterable)
while True:
values = []
for _ in range(n):
try:
value = next(it)
except StopIteration:
values.extend([fillvalue] * (n - len(values)))
yield tuple(values)
return
values.append(value)
yield tuple(values)
class TestCombinedDataLoader(unittest.TestCase):
def test_combine_loaders_1(self):
loader1 = _grouper([f"1_{i}" for i in range(10)], 2)
loader2 = _grouper([f"2_{i}" for i in range(11)], 3)
batch_size = 4
ratios = (0.1, 0.9)
random.seed(43)
combined = CombinedDataLoader((loader1, loader2), batch_size, ratios)
BATCHES_GT = [
["1_0", "1_1", "2_0", "2_1"],
["2_2", "2_3", "2_4", "2_5"],
["1_2", "1_3", "2_6", "2_7"],
["2_8", "2_9", "2_10", None],
]
for i, batch in enumerate(combined):
self.assertEqual(len(batch), batch_size)
self.assertEqual(batch, BATCHES_GT[i])
|