添加了 reproducible batch sampler 的测试

fastNLP/core/samplers/reproducible_batch_sampler.py

@@ -8,7 +8,6 @@ import math
 from copy import deepcopy
 from typing import Dict, Union, List
 from itertools import chain
-import os
 
 import numpy as np
 

tests/core/controllers/utils/test_utils.py

@@ -1,7 +1,7 @@
 from functools import reduce
 
 from fastNLP.core.controllers.utils.utils import _TruncatedDataLoader  # TODO: 该类修改过，记得将 test 也修改；
-from tests.helpers.datasets.normal_data import NormalIterator
+from tests.helpers.datasets.normal_data import NormalSampler
 
 
 class Test_WrapDataLoader:
@@ -9,7 +9,7 @@ class Test_WrapDataLoader:
     def test_normal_generator(self):
         all_sanity_batches = [4, 20, 100]
         for sanity_batches in all_sanity_batches:
-            data = NormalIterator(num_of_data=1000)
+            data = NormalSampler(num_of_data=1000)
             wrapper = _TruncatedDataLoader(dataloader=data, num_batches=sanity_batches)
             dataloader = iter(wrapper(dataloader=data))
             mark = 0

tests/core/samplers/test_reproducible_batch_sampler.py

@@ -1,161 +1,131 @@
-from array import array
-
 import numpy as np
 import pytest
 from itertools import chain
 from copy import deepcopy
+from array import array
 
+from tests.helpers.datasets.normal_data import NormalSampler, NormalBatchSampler
 from fastNLP.core.samplers import ReproduceBatchSampler, BucketedBatchSampler, RandomBatchSampler
-from fastNLP.core.drivers.torch_driver.utils import replace_batch_sampler
-from tests.helpers.datasets.torch_data import TorchNormalDataset
 
-#
-# class TestReproducibleBatchSampler:
-#     # TODO 拆分测试，在这里只测试一个东西
-#     def test_torch_dataloader_1(self):
-#         import torch
-#         from torch.utils.data import DataLoader
-#         # no shuffle
-#         before_batch_size = 7
-#         dataset = TorchNormalDataset(num_of_data=100)
-#         dataloader = DataLoader(dataset, batch_size=before_batch_size)
-#         re_batchsampler = ReproduceBatchSampler(dataloader.batch_sampler, dataloader.batch_size, drop_last=False)
-#         dataloader = replace_batch_sampler(dataloader, re_batchsampler)
-#
-#         forward_steps = 3
-#         iter_dataloader = iter(dataloader)
-#         for _ in range(forward_steps):
-#             next(iter_dataloader)
-#
-#         # 1. 保存状态
-#         _get_re_batchsampler = dataloader.batch_sampler
-#         assert isinstance(_get_re_batchsampler, ReproduceBatchSampler)
-#         state = _get_re_batchsampler.state_dict()
-#         assert state == {"index_list": array("I", list(range(100))), "num_consumed_samples": forward_steps*before_batch_size,
-#                          "sampler_type": "ReproduceBatchSampler"}
-#
-#         # 2. 断点重训，重新生成一个 dataloader；
-#         # 不改变 batch_size；
-#         dataloader = DataLoader(dataset, batch_size=before_batch_size)
-#         re_batchsampler = ReproduceBatchSampler(dataloader.batch_sampler, dataloader.batch_size, drop_last=False)
-#         re_batchsampler.load_state_dict(state)
-#         dataloader = replace_batch_sampler(dataloader, re_batchsampler)
-#
-#         real_res = []
-#         supposed_res = (torch.tensor(list(range(21, 28))), torch.tensor(list(range(28, 35))))
-#         forward_steps = 2
-#         iter_dataloader = iter(dataloader)
-#         for _ in range(forward_steps):
-#             real_res.append(next(iter_dataloader))
-#
-#         for i in range(forward_steps):
-#             assert all(real_res[i] == supposed_res[i])
-#
-#         # 改变 batch_size；
-#         after_batch_size = 3
-#         dataloader = DataLoader(dataset, batch_size=after_batch_size)
-#         re_batchsampler = ReproduceBatchSampler(dataloader.batch_sampler, dataloader.batch_size, drop_last=False)
-#         re_batchsampler.load_state_dict(state)
-#         dataloader = replace_batch_sampler(dataloader, re_batchsampler)
-#
-#         real_res = []
-#         supposed_res = (torch.tensor(list(range(21, 24))), torch.tensor(list(range(24, 27))))
-#         forward_steps = 2
-#         iter_dataloader = iter(dataloader)
-#         for _ in range(forward_steps):
-#             real_res.append(next(iter_dataloader))
-#
-#         for i in range(forward_steps):
-#             assert all(real_res[i] == supposed_res[i])
-#
-#         # 断点重训的第二轮是否是一个完整的 dataloader；
-#         # 先把断点重训所在的那一个 epoch 跑完；
-#         begin_idx = 27
-#         while True:
-#             try:
-#                 data = next(iter_dataloader)
-#                 _batch_size = len(data)
-#                 assert all(data == torch.tensor(list(range(begin_idx, begin_idx + _batch_size))))
-#                 begin_idx += _batch_size
-#             except StopIteration:
-#                 break
-#
-#         # 开始新的一轮；
-#         begin_idx = 0
-#         iter_dataloader = iter(dataloader)
-#         while True:
-#             try:
-#                 data = next(iter_dataloader)
-#                 _batch_size = len(data)
-#                 assert all(data == torch.tensor(list(range(begin_idx, begin_idx + _batch_size))))
-#                 begin_idx += _batch_size
-#             except StopIteration:
-#                 break
-#
-#     def test_torch_dataloader_2(self):
-#         # 测试新的一轮的 index list 是重新生成的，而不是沿用上一轮的；
-#         from torch.utils.data import DataLoader
-#         # no shuffle
-#         before_batch_size = 7
-#         dataset = TorchNormalDataset(num_of_data=100)
-#         # 开启 shuffle，来检验断点重训后的第二轮的 index list 是不是重新生成的；
-#         dataloader = DataLoader(dataset, batch_size=before_batch_size, shuffle=True)
-#         re_batchsampler = ReproduceBatchSampler(dataloader.batch_sampler, dataloader.batch_size, drop_last=False)
-#         dataloader = replace_batch_sampler(dataloader, re_batchsampler)
-#
-#         # 将一轮的所有数据保存下来，看是否恢复的是正确的；
-#         all_supposed_data = []
-#         forward_steps = 3
-#         iter_dataloader = iter(dataloader)
-#         for _ in range(forward_steps):
-#             all_supposed_data.extend(next(iter_dataloader).tolist())
-#
-#         # 1. 保存状态
-#         _get_re_batchsampler = dataloader.batch_sampler
-#         assert isinstance(_get_re_batchsampler, ReproduceBatchSampler)
-#         state = _get_re_batchsampler.state_dict()
-#
-#         # 2. 断点重训，重新生成一个 dataloader；
-#         # 不改变 batch_size；
-#         dataloader = DataLoader(dataset, batch_size=before_batch_size, shuffle=True)
-#         re_batchsampler = ReproduceBatchSampler(dataloader.batch_sampler, dataloader.batch_size, drop_last=False)
-#         re_batchsampler.load_state_dict(state)
-#         dataloader = replace_batch_sampler(dataloader, re_batchsampler)
-#
-#         # 先把这一轮的数据过完；
-#         pre_index_list = dataloader.batch_sampler.state_dict()["index_list"]
-#         while True:
-#             try:
-#                 all_supposed_data.extend(next(iter_dataloader).tolist())
-#             except StopIteration:
-#                 break
-#         assert all_supposed_data == list(pre_index_list)
-#
-#         # 重新开启新的一轮；
-#         for _ in range(3):
-#             iter_dataloader = iter(dataloader)
-#             res = []
-#             while True:
-#                 try:
-#                     res.append(next(iter_dataloader))
-#                 except StopIteration:
-#                     break
-#
-#     def test_3(self):
-#         import torch
-#         from torch.utils.data import DataLoader
-#         before_batch_size = 7
-#         dataset = TorchNormalDataset(num_of_data=100)
-#         # 开启 shuffle，来检验断点重训后的第二轮的 index list 是不是重新生成的；
-#         dataloader = DataLoader(dataset, batch_size=before_batch_size)
-#
-#         for idx, data in enumerate(dataloader):
-#             if idx > 3:
-#                 break
-#
-#         iterator = iter(dataloader)
-#         for each in iterator:
-#             pass
+
+class TestReproducibleBatchSampler:
+    def test_1(self):
+        sampler = NormalSampler(num_of_data=100)  # 这里是否是 batchsampler 不影响；
+
+        reproduce_batch_sampler = ReproduceBatchSampler(sampler, batch_size=4, drop_last=False)
+
+        forward_steps = 3
+        iterator = iter(reproduce_batch_sampler)
+        i = 0
+        while i < forward_steps:
+            next(iterator)
+            i += 1
+
+        # 保存状态；
+        state = reproduce_batch_sampler.state_dict()
+
+        assert state == {"index_list": array("I", list(range(100))),
+                         "num_consumed_samples": forward_steps * 4,
+                         "sampler_type": "ReproduceBatchSampler"}
+
+        # 重新生成一个 batchsampler 然后加载状态；
+        sampler = NormalSampler(num_of_data=100)  # 这里是否是 batchsampler 不影响；
+        reproduce_batch_sampler = ReproduceBatchSampler(sampler, batch_size=4, drop_last=False)
+        reproduce_batch_sampler.load_state_dict(state)
+
+        real_res = []
+        supposed_res = (list(range(12, 16)), list(range(16, 20)))
+        forward_steps = 2
+        iter_dataloader = iter(reproduce_batch_sampler)
+        for _ in range(forward_steps):
+            real_res.append(next(iter_dataloader))
+
+        for i in range(forward_steps):
+            assert supposed_res[i] == real_res[i]
+
+        # 改变 batchsize；
+        sampler = NormalSampler(num_of_data=100)  # 这里是否是 batchsampler 不影响；
+        reproduce_batch_sampler = ReproduceBatchSampler(sampler, batch_size=7, drop_last=False)
+        reproduce_batch_sampler.load_state_dict(state)
+
+        real_res = []
+        supposed_res = (list(range(12, 19)), list(range(19, 26)))
+        forward_steps = 2
+        iter_dataloader = iter(reproduce_batch_sampler)
+        for _ in range(forward_steps):
+            real_res.append(next(iter_dataloader))
+
+        for i in range(forward_steps):
+            assert supposed_res[i] == real_res[i]
+
+        # 断点重训的第二轮是否是一个完整的 dataloader；
+        # 先把断点重训所在的那一个 epoch 跑完；
+        begin_idx = 26
+        while True:
+            try:
+                data = next(iter_dataloader)
+                _batch_size = len(data)
+                assert data == list(range(begin_idx, begin_idx + _batch_size))
+                begin_idx += _batch_size
+            except StopIteration:
+                break
+
+        # 开始新的一轮；
+        begin_idx = 0
+        iter_dataloader = iter(reproduce_batch_sampler)
+        while True:
+            try:
+                data = next(iter_dataloader)
+                _batch_size = len(data)
+                assert data == list(range(begin_idx, begin_idx + _batch_size))
+                begin_idx += _batch_size
+            except StopIteration:
+                break
+
+    def test_2(self):
+
+        # 测试新的一轮的 index list 是重新生成的，而不是沿用上一轮的；
+        before_batch_size = 7
+        sampler = NormalSampler(num_of_data=100)
+        # 开启 shuffle，来检验断点重训后的第二轮的 index list 是不是重新生成的；
+        reproduce_batch_sampler = ReproduceBatchSampler(sampler, before_batch_size, drop_last=False)
+
+        # 将一轮的所有数据保存下来，看是否恢复的是正确的；
+        all_supposed_data = []
+        forward_steps = 3
+        iter_dataloader = iter(reproduce_batch_sampler)
+        for _ in range(forward_steps):
+            all_supposed_data.extend(next(iter_dataloader))
+
+        # 1. 保存状态
+        state = reproduce_batch_sampler.state_dict()
+
+        # 2. 断点重训，重新生成一个 dataloader；
+        # 不改变 batch_size；
+        sampler = NormalSampler(num_of_data=100, shuffle=True)
+        reproduce_batch_sampler = ReproduceBatchSampler(sampler, before_batch_size, drop_last=False)
+        reproduce_batch_sampler.load_state_dict(state)
+
+        # 先把这一轮的数据过完；
+        pre_index_list = reproduce_batch_sampler.state_dict()["index_list"]
+        iter_dataloader = iter(reproduce_batch_sampler)
+        while True:
+            try:
+                all_supposed_data.extend(next(iter_dataloader))
+            except StopIteration:
+                break
+        assert all_supposed_data == list(pre_index_list)
+
+        # 重新开启新的一轮；
+        for _ in range(3):
+            iter_dataloader = iter(reproduce_batch_sampler)
+            res = []
+            while True:
+                try:
+                    res.extend(next(iter_dataloader))
+                except StopIteration:
+                    break
+            assert res != all_supposed_data
 
 
 class DatasetWithVaryLength:

tests/core/samplers/test_reproducible_batch_sampler_torch.py

@@ -0,0 +1,141 @@
+from array import array
+import torch
+from torch.utils.data import DataLoader
+
+import pytest
+
+from fastNLP.core.samplers import ReproduceBatchSampler
+from fastNLP.core.drivers.torch_driver.utils import replace_batch_sampler
+from tests.helpers.datasets.torch_data import TorchNormalDataset
+
+
+@pytest.mark.torch
+class TestReproducibleBatchSamplerTorch:
+    def test_torch_dataloader_1(self):
+        # no shuffle
+        before_batch_size = 7
+        dataset = TorchNormalDataset(num_of_data=100)
+        dataloader = DataLoader(dataset, batch_size=before_batch_size)
+        re_batchsampler = ReproduceBatchSampler(dataloader.batch_sampler, dataloader.batch_size, drop_last=False)
+        dataloader = replace_batch_sampler(dataloader, re_batchsampler)
+
+        forward_steps = 3
+        iter_dataloader = iter(dataloader)
+        for _ in range(forward_steps):
+            next(iter_dataloader)
+
+        # 1. 保存状态
+        _get_re_batchsampler = dataloader.batch_sampler
+        assert isinstance(_get_re_batchsampler, ReproduceBatchSampler)
+        state = _get_re_batchsampler.state_dict()
+        assert state == {"index_list": array("I", list(range(100))), "num_consumed_samples": forward_steps*before_batch_size,
+                         "sampler_type": "ReproduceBatchSampler"}
+
+        # 2. 断点重训，重新生成一个 dataloader；
+        # 不改变 batch_size；
+        dataloader = DataLoader(dataset, batch_size=before_batch_size)
+        re_batchsampler = ReproduceBatchSampler(dataloader.batch_sampler, dataloader.batch_size, drop_last=False)
+        re_batchsampler.load_state_dict(state)
+        dataloader = replace_batch_sampler(dataloader, re_batchsampler)
+
+        real_res = []
+        supposed_res = (torch.tensor(list(range(21, 28))), torch.tensor(list(range(28, 35))))
+        forward_steps = 2
+        iter_dataloader = iter(dataloader)
+        for _ in range(forward_steps):
+            real_res.append(next(iter_dataloader))
+
+        for i in range(forward_steps):
+            assert all(real_res[i] == supposed_res[i])
+
+        # 改变 batch_size；
+        after_batch_size = 3
+        dataloader = DataLoader(dataset, batch_size=after_batch_size)
+        re_batchsampler = ReproduceBatchSampler(dataloader.batch_sampler, dataloader.batch_size, drop_last=False)
+        re_batchsampler.load_state_dict(state)
+        dataloader = replace_batch_sampler(dataloader, re_batchsampler)
+
+        real_res = []
+        supposed_res = (torch.tensor(list(range(21, 24))), torch.tensor(list(range(24, 27))))
+        forward_steps = 2
+        iter_dataloader = iter(dataloader)
+        for _ in range(forward_steps):
+            real_res.append(next(iter_dataloader))
+
+        for i in range(forward_steps):
+            assert all(real_res[i] == supposed_res[i])
+
+        # 断点重训的第二轮是否是一个完整的 dataloader；
+        # 先把断点重训所在的那一个 epoch 跑完；
+        begin_idx = 27
+        while True:
+            try:
+                data = next(iter_dataloader)
+                _batch_size = len(data)
+                assert all(data == torch.tensor(list(range(begin_idx, begin_idx + _batch_size))))
+                begin_idx += _batch_size
+            except StopIteration:
+                break
+
+        # 开始新的一轮；
+        begin_idx = 0
+        iter_dataloader = iter(dataloader)
+        while True:
+            try:
+                data = next(iter_dataloader)
+                _batch_size = len(data)
+                assert all(data == torch.tensor(list(range(begin_idx, begin_idx + _batch_size))))
+                begin_idx += _batch_size
+            except StopIteration:
+                break
+
+    def test_torch_dataloader_2(self):
+        # 测试新的一轮的 index list 是重新生成的，而不是沿用上一轮的；
+        from torch.utils.data import DataLoader
+        before_batch_size = 7
+        dataset = TorchNormalDataset(num_of_data=100)
+        # 开启 shuffle，来检验断点重训后的第二轮的 index list 是不是重新生成的；
+        dataloader = DataLoader(dataset, batch_size=before_batch_size, shuffle=True)
+        re_batchsampler = ReproduceBatchSampler(dataloader.batch_sampler, dataloader.batch_size, drop_last=False)
+        dataloader = replace_batch_sampler(dataloader, re_batchsampler)
+
+        # 将一轮的所有数据保存下来，看是否恢复的是正确的；
+        all_supposed_data = []
+        forward_steps = 3
+        iter_dataloader = iter(dataloader)
+        for _ in range(forward_steps):
+            all_supposed_data.extend(next(iter_dataloader).tolist())
+
+        # 1. 保存状态
+        _get_re_batchsampler = dataloader.batch_sampler
+        assert isinstance(_get_re_batchsampler, ReproduceBatchSampler)
+        state = _get_re_batchsampler.state_dict()
+
+        # 2. 断点重训，重新生成一个 dataloader；
+        # 不改变 batch_size；
+        dataloader = DataLoader(dataset, batch_size=before_batch_size, shuffle=True)
+        re_batchsampler = ReproduceBatchSampler(dataloader.batch_sampler, dataloader.batch_size, drop_last=False)
+        re_batchsampler.load_state_dict(state)
+        dataloader = replace_batch_sampler(dataloader, re_batchsampler)
+
+        iter_dataloader = iter(dataloader)
+        # 先把这一轮的数据过完；
+        pre_index_list = dataloader.batch_sampler.state_dict()["index_list"]
+        while True:
+            try:
+                all_supposed_data.extend(next(iter_dataloader).tolist())
+            except StopIteration:
+                break
+        assert all_supposed_data == list(pre_index_list)
+
+        # 重新开启新的一轮；
+        for _ in range(3):
+            iter_dataloader = iter(dataloader)
+            res = []
+            while True:
+                try:
+                    res.extend(next(iter_dataloader).tolist())
+                except StopIteration:
+                    break
+            assert res != all_supposed_data
+

tests/helpers/datasets/normal_data.py

@@ -1,13 +1,25 @@
 import numpy as np
+import random
 
 
-class NormalIterator:
-    def __init__(self, num_of_data=1000):
+class NormalSampler:
+    def __init__(self, num_of_data=1000, shuffle=False):
         self._num_of_data = num_of_data
         self._data = list(range(num_of_data))
+        if shuffle:
+            random.shuffle(self._data)
+        self.shuffle = shuffle
         self._index = 0
+        self.need_reinitialize = False
 
     def __iter__(self):
+        if self.need_reinitialize:
+            self._index = 0
+            if self.shuffle:
+                random.shuffle(self._data)
+        else:
+            self.need_reinitialize = True
+
         return self
 
     def __next__(self):
@@ -15,12 +27,45 @@ class NormalIterator:
             raise StopIteration
         _data = self._data[self._index]
         self._index += 1
-        return self._data
+        return _data
 
     def __len__(self):
         return self._num_of_data
 
 
+class NormalBatchSampler:
+    def __init__(self, sampler, batch_size: int, drop_last: bool) -> None:
+        # Since collections.abc.Iterable does not check for `__getitem__`, which
+        # is one way for an object to be an iterable, we don't do an `isinstance`
+        # check here.
+        if not isinstance(batch_size, int) or isinstance(batch_size, bool) or \
+                batch_size <= 0:
+            raise ValueError("batch_size should be a positive integer value, "
+                             "but got batch_size={}".format(batch_size))
+        if not isinstance(drop_last, bool):
+            raise ValueError("drop_last should be a boolean value, but got "
+                             "drop_last={}".format(drop_last))
+        self.sampler = sampler
+        self.batch_size = batch_size
+        self.drop_last = drop_last
+
+    def __iter__(self):
+        batch = []
+        for idx in self.sampler:
+            batch.append(idx)
+            if len(batch) == self.batch_size:
+                yield batch
+                batch = []
+        if len(batch) > 0 and not self.drop_last:
+            yield batch
+
+    def __len__(self) -> int:
+        if self.drop_last:
+            return len(self.sampler) // self.batch_size
+        else:
+            return (len(self.sampler) + self.batch_size - 1) // self.batch_size
+
+
 class RandomDataset:
     def __init__(self, num_data=10):
         self.data = np.random.rand(num_data)
@@ -30,3 +75,6 @@ class RandomDataset:
 
     def __getitem__(self, item):
         return self.data[item]
+
+
+