Change DummySerializer to support BalancedSerialIterator
@@ -4,23 +4,22 @@ import pytest
from chainer import serializer
from chainer_chemistry.datasets.numpy_tuple_dataset import NumpyTupleDataset
-
from chainer_chemistry.iterators.balanced_serial_iterator import BalancedSerialIterator # NOQA
from chainer_chemistry.iterators.balanced_serial_iterator import IndexIterator # NOQA
class DummySerializer(serializer.Serializer):
- def __init__(self, target):
- super(DummySerializer, self).__init__()
- self.target = target
+ def __init__(self, path=[]):
+ self.path = path
+ self.called = []
def __getitem__(self, key):
- raise NotImplementedError
+ return DummySerializer(self.path + [key])
def __call__(self, key, value):
- self.target[key] = value
- return self.target[key]
+ self.called.append((key, value))
+ return value
def test_index_iterator():
@@ -55,8 +54,7 @@ def _test_index_iterator_no_shuffle():
assert indices3[0] == index_list[1]
assert indices3[1] == index_list[2]
- target = dict()
- ii.serialize(DummySerializer(target))
+ ii.serialize(DummySerializer())
assert isinstance(ii.current_index_list, numpy.ndarray)
assert len(ii.current_index_list) == len(index_list)
assert numpy.array_equal(ii.current_index_list, numpy.asarray(index_list))
@@ -82,8 +80,7 @@ def _test_index_iterator_with_shuffle():
for index in indices:
assert index in index_list
- target = dict()
- ii.serialize(DummySerializer(target))
+ ii.serialize(DummySerializer())
for index in ii.current_index_list:
assert index in index_list
assert ii.current_pos == (3 + 6) % len(index_list) + 2
@@ -118,9 +115,13 @@ def _test_balanced_serial_iterator_no_batch_balancing():
assert numpy.sum(labels_batch == 1) == 3
assert numpy.sum(labels_batch == 2) == 3
- # This does not work!
- # target = dict()
- # iterator.serialize(DummySerializer(target))
+ iterator.serialize(DummySerializer())
+ assert iterator.current_position == 0
+ assert iterator.epoch == 1
+ assert iterator.is_new_epoch == True
+ for order in iterator._order:
+ assert order <= 9
+ assert iterator._previous_epoch_detail == 0.0
def _test_balanced_serial_iterator_with_batch_balancing():
@@ -143,9 +144,12 @@ def _test_balanced_serial_iterator_with_batch_balancing():
assert numpy.sum(labels_batch == 1) == 1
assert numpy.sum(labels_batch == 2) == 1
- # This does not work!
- # target = dict()
- # iterator.serialize(DummySerializer(target))
+ iterator.serialize(DummySerializer())
+ assert iterator.current_position == 0
+ assert iterator.epoch == 1
+ assert iterator.is_new_epoch == True
+ for order in iterator._order:
+ assert order <= 9
if __name__ == '__main__':
