-
Notifications
You must be signed in to change notification settings - Fork 244
Expand file tree
/
Copy pathtest_batching.py
More file actions
46 lines (39 loc) · 1.25 KB
/
Copy pathtest_batching.py
File metadata and controls
46 lines (39 loc) · 1.25 KB
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
import numpy as np
import pytest
from docarray import BaseDoc, DocArray
from docarray.typing import NdArray
@pytest.mark.parametrize('shuffle', [False, True])
@pytest.mark.parametrize('stack', [False, True])
@pytest.mark.parametrize('batch_size,n_batches', [(16, 7), (10, 10)])
def test_batch(shuffle, stack, batch_size, n_batches):
class MyDoc(BaseDoc):
id: int
tensor: NdArray
t_shape = (32, 32)
da = DocArray[MyDoc](
[
MyDoc(
id=i,
tensor=np.zeros(t_shape),
)
for i in range(100)
]
)
if stack:
da = da.stack()
batches = list(da._batch(batch_size=batch_size, shuffle=shuffle))
assert len(batches) == n_batches
for i, batch in enumerate(batches):
if i < n_batches - 1:
assert len(batch) == batch_size
if stack:
assert batch.tensor.shape == (batch_size, *t_shape)
else:
assert len(batch) <= batch_size
non_shuffled_ids = [
i for i in range(i * batch_size, min((i + 1) * batch_size, len(da)))
]
if not shuffle:
assert batch.id == non_shuffled_ids
else:
assert not (batch.id == non_shuffled_ids)