|
- from functools import reduce
-
- from fastNLP.core.controllers.utils.utils import _TruncatedDataLoader # TODO: 该类修改过，记得将 test 也修改；
- from tests.helpers.datasets.normal_data import NormalIterator
-
-
- 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)
- wrapper = _TruncatedDataLoader(num_batches=sanity_batches)
- dataloader = iter(wrapper(dataloader=data))
- mark = 0
- while True:
- try:
- _data = next(dataloader)
- except StopIteration:
- break
- mark += 1
- assert mark == sanity_batches
-
- def test_torch_dataloader(self):
- from tests.helpers.datasets.torch_data import TorchNormalDataset
- from torch.utils.data import DataLoader
-
- bses = [8, 16, 40]
- all_sanity_batches = [4, 7, 10]
- for bs in bses:
- for sanity_batches in all_sanity_batches:
- dataset = TorchNormalDataset(num_of_data=1000)
- dataloader = DataLoader(dataset, batch_size=bs, shuffle=True)
- wrapper = _TruncatedDataLoader(num_batches=sanity_batches)
- dataloader = wrapper(dataloader)
- dataloader = iter(dataloader)
- all_supposed_running_data_num = 0
- while True:
- try:
- _data = next(dataloader)
- except StopIteration:
- break
- all_supposed_running_data_num += _data.shape[0]
- assert all_supposed_running_data_num == bs * sanity_batches
-
- def test_len(self):
- from tests.helpers.datasets.torch_data import TorchNormalDataset
- from torch.utils.data import DataLoader
-
- bses = [8, 16, 40]
- all_sanity_batches = [4, 7, 10]
- length = []
- for bs in bses:
- for sanity_batches in all_sanity_batches:
- dataset = TorchNormalDataset(num_of_data=1000)
- dataloader = DataLoader(dataset, batch_size=bs, shuffle=True)
- wrapper = _TruncatedDataLoader(num_batches=sanity_batches)
- dataloader = wrapper(dataloader)
- length.append(len(dataloader))
- assert length == reduce(lambda x, y: x+y, [all_sanity_batches for _ in range(len(bses))])
|
