import unittest from itertools import product import numpy as np from functools import partial from array import array from fastNLP.core.samplers.reproducible_sampler import RandomSampler from fastNLP.core.drivers.torch_driver.utils import replace_batch_sampler from tests.helpers.datasets.torch_data import TorchNormalDataset class TestRandomSamplerYh(unittest.TestCase): def test_init(self): # 测试能否正确初始化 dataset = TorchNormalDataset(num_of_data=100) sampler = RandomSampler(dataset) for i in sampler: pass def test_during_iter(self): dataset = TorchNormalDataset(num_of_data=100) sampler = RandomSampler(dataset) for i in sampler: with self.assertRaises(AssertionError): sampler.set_distributed(1, 0) break # should not raise for i in sampler: pass sampler.set_distributed(1, 0) def test_set_distributed(self): dataset = TorchNormalDataset(num_of_data=100) sampler = RandomSampler(dataset, shuffle=False) sampler.set_distributed(num_replicas=2, rank=0, pad=False) self.assertEqual(len(sampler), 50) count = 0 for i in sampler: self.assertEqual(i%2, 0) count += 1 self.assertEqual(count, 50) sampler.set_distributed(num_replicas=2, rank=1, pad=False) self.assertEqual(len(sampler), 50) count = 0 for i in sampler: self.assertEqual(i%2, 1) count += 1 self.assertEqual(count, 50) dataset = TorchNormalDataset(num_of_data=101) sampler = RandomSampler(dataset, shuffle=False) sampler.set_distributed(num_replicas=2, rank=0, pad=True) self.assertEqual(len(sampler), 51) count = 0 for i in sampler: self.assertEqual(i%2, 0) count += 1 self.assertEqual(count, 51) sampler.set_distributed(num_replicas=2, rank=1, pad=True) self.assertEqual(len(sampler), 51) count = 0 for i in sampler: if i!=0: self.assertEqual(i%2, 1) count += 1 self.assertEqual(count, 51) def test_state_dict_check_length(self): dataset = TorchNormalDataset(num_of_data=100) sampler = RandomSampler(dataset, shuffle=False) states = sampler.state_dict() new_ds = TorchNormalDataset(num_of_data=10) with self.assertRaises(AssertionError): new_sampler = RandomSampler(new_ds) new_sampler.load_state_dict(states) new_ds = TorchNormalDataset(num_of_data=100) new_sampler = RandomSampler(new_ds) new_sampler.load_state_dict(states) def test_state_dict(self): num_samples = 100 dataset = TorchNormalDataset(num_of_data=num_samples) # 测试使用 前后shuffle不一致的load操作 lst = [0]+np.random.randint(1, num_samples, size=3).tolist() for pre_shuffle, post_shuffle, num_consumed_samples in product([True, False], [True, False], lst): with self.subTest(pre_shuffle=pre_shuffle, post_shuffle=post_shuffle, num_consumed_samples=num_consumed_samples): sampler = RandomSampler(dataset, shuffle=pre_shuffle) sampler.set_epoch(0) already_numbers = set() if num_consumed_samples>0: for i, j in enumerate(sampler, start=1): already_numbers.add(j) if i == num_consumed_samples: break self.assertEqual(len(already_numbers), num_consumed_samples) states = sampler.state_dict() new_sampler = RandomSampler(dataset, shuffle=post_shuffle) new_sampler.load_state_dict(states) new_sampler.set_epoch(0) for i in new_sampler: self.assertNotIn(i, already_numbers) # 测试切换成多卡也没有问题 other_rank_number = set() for rank in range(3): new_sampler = RandomSampler(dataset, shuffle=post_shuffle) new_sampler.load_state_dict(states) new_sampler.set_distributed(num_replicas=3, rank=rank, pad=False) new_sampler.set_epoch(0) count = 0 for i in new_sampler: self.assertNotIn(i, other_rank_number) other_rank_number.add(i) self.assertNotIn(i, already_numbers) count += 1 def test_state_dict_2(self): # 测试一下从多卡切换到单卡，或者切换到不同卡数量的多卡 num_samples = 100 dataset = TorchNormalDataset(num_of_data=num_samples) # 测试使用 前后shuffle不一致的load操作 lst = [0]+np.random.randint(1, num_samples//2, size=3).tolist() # lst = [30] for pre_shuffle, post_shuffle, num_consumed_samples in product([True, False], [True, False], lst): with self.subTest(pre_shuffle=pre_shuffle, post_shuffle=post_shuffle, num_consumed_samples=num_consumed_samples): already_numbers = set() sampler = RandomSampler(dataset, shuffle=pre_shuffle, seed=0) sampler.set_distributed(num_replicas=2, rank=0) sampler.set_epoch(0) if num_consumed_samples>0: for i, j in enumerate(sampler, start=1): already_numbers.add(j) if i == num_consumed_samples: break sampler = RandomSampler(dataset, shuffle=pre_shuffle, seed=0) sampler.set_epoch(0) sampler.set_distributed(num_replicas=2, rank=1) if num_consumed_samples>0: for i, j in enumerate(sampler, start=1): already_numbers.add(j) if i == num_consumed_samples: break self.assertEqual(len(already_numbers), num_consumed_samples*2) states = sampler.state_dict() new_sampler = RandomSampler(dataset, shuffle=post_shuffle) new_sampler.load_state_dict(states) new_sampler.set_epoch(0) for i in new_sampler: self.assertNotIn(i, already_numbers) # 测试切换成多卡也没有问题 other_rank_number = set() for rank in range(3): new_sampler = RandomSampler(dataset, shuffle=post_shuffle) new_sampler.load_state_dict(states) new_sampler.set_epoch(0) new_sampler.set_distributed(num_replicas=3, rank=rank, pad=False) count = 0 for i in new_sampler: self.assertNotIn(i, other_rank_number) other_rank_number.add(i) self.assertNotIn(i, already_numbers) count += 1 class TestRandomSampler(unittest.TestCase): # 测试单卡； def test_seed_work_when_shuffle_is_true(self): data_length = 100 torch_normal_data = TorchNormalDataset(num_of_data=data_length) for shuffle in [True, False]: iterable = RandomSampler(dataset=torch_normal_data, shuffle=shuffle) # 迭代一些数据，但是不迭代完； iterable.set_epoch(1) iterator = iter(iterable) pre_data = [] forward_steps = 30 for _ in range(forward_steps): pre_data.append(next(iterator)) # 看重新生成迭代器是否能够完全重置状态； iterator = iter(iterable) res = [] for _ in range(forward_steps): res.append(next(iterator)) assert pre_data == res # 测试断点重训； # 如果 shuffle，那么下一轮的数据应当与前一轮不一样；并且如果是断点重训，两次的下一轮应当是一样的； def test_2(self): data_length = 100 torch_normal_data = TorchNormalDataset(num_of_data=data_length) random_sampler_1 = RandomSampler(dataset=torch_normal_data, shuffle=True) iterator = iter(random_sampler_1) # 第一轮 random_sampler_1.set_epoch(0) first_epoch = [] forward_steps = 30 for _ in range(forward_steps): first_epoch.append(next(iterator)) # 先提前保存断点重训的结果; state = random_sampler_1.state_dict() # 保存第一个 epoch 的之后的结果，用于查看断点重训是否正确； first_left_data = [] while True: try: first_left_data.append(next(iterator)) except StopIteration: break # 第二轮 random_sampler_1.set_epoch(1) iterator = iter(random_sampler_1) second_epoch = [] for _ in range(forward_steps): second_epoch.append(next(iterator)) assert first_epoch != second_epoch # 重新加载第一轮的状态，查看断点重训是否正确； random_sampler_2 = RandomSampler(dataset=torch_normal_data, shuffle=True) random_sampler_2.load_state_dict(state) random_sampler_2.set_epoch(0) iterator = iter(random_sampler_2) re_first_epoch = [] while True: try: re_first_epoch.append(next(iterator)) except StopIteration: break assert re_first_epoch == first_left_data # 查看第二轮的结果是否也是和第一次的第二轮完全一致； random_sampler_2.set_epoch(1) iterator = iter(random_sampler_2) re_second_epoch = [] for _ in range(forward_steps): re_second_epoch.append(next(iterator)) assert re_second_epoch == second_epoch # 多卡； # 如果一个 sampler 还没有迭代完，我们又直接 iter(sampler) 那么是否正确（应当生成一个全新的 sampler）？ def test_3(self): data_length = 100 torch_normal_data = TorchNormalDataset(num_of_data=data_length) random_sampler_1 = partial(RandomSampler, dataset=torch_normal_data, shuffle=False) random_sampler_2 = partial(RandomSampler, dataset=torch_normal_data, shuffle=True) iterable_items = [random_sampler_1, random_sampler_2] world_size = 3 for pad in {True, False}: for iterable in iterable_items: for rank in range(world_size): each_rank_iterable = iterable() each_rank_iterable.set_epoch(0) each_rank_iterable.set_distributed(num_replicas=world_size, rank=rank, pad=pad) # 迭代一些数据，但是不迭代完； iterator = iter(each_rank_iterable) pre_data = [] forward_steps = 10 for _ in range(forward_steps): pre_data.append(next(iterator)) # 看重新生成迭代器是否能够完全重置状态； iterator = iter(each_rank_iterable) res = [] for _ in range(forward_steps): res.append(next(iterator)) assert res == pre_data # 测试断点重训； # 如果 shuffle，那么下一轮的数据应当与前一轮不一样；并且如果是断点重训，两次的下一轮应当是一样的； def test_4(self): data_length = 100 torch_normal_data = TorchNormalDataset(num_of_data=data_length) random_sampler_1 = partial(RandomSampler, dataset=torch_normal_data, shuffle=True) world_size_1 = 2 forward_steps = 10 for pad in {True, False}: all_rank_state = {} all_rank_first_left_data = {} all_rank_second_epoch = {} for rank in range(world_size_1): each_rank_iterable = random_sampler_1() each_rank_iterable.set_distributed(num_replicas=world_size_1, rank=rank, pad=pad) iterator = iter(each_rank_iterable) # 第一轮 each_rank_iterable.set_epoch(0) first_epoch = [] for _ in range(forward_steps): first_epoch.append(next(iterator)) # 先提前保存断点重训的结果; all_rank_state[rank] = each_rank_iterable.state_dict() # 保存第一个 epoch 的之后的结果，用于查看断点重训是否正确； first_left_data = [] while True: try: first_left_data.append(next(iterator)) except StopIteration: break all_rank_first_left_data[rank] = first_left_data # 第二轮 each_rank_iterable.set_epoch(1) iterator = iter(each_rank_iterable) second_epoch = [] for _ in range(forward_steps): second_epoch.append(next(iterator)) all_rank_second_epoch[rank] = second_epoch assert first_epoch != second_epoch # 重新加载第一轮的状态，查看断点重训是否正确； random_sampler_2 = partial(RandomSampler, dataset=torch_normal_data, shuffle=True) for rank in range(world_size_1): each_rank_iterable = random_sampler_2() each_rank_iterable.set_distributed(num_replicas=world_size_1, rank=rank, pad=pad) each_rank_iterable.load_state_dict(all_rank_state[rank]) each_rank_iterable.set_epoch(0) iterator = iter(each_rank_iterable) re_first_epoch = [] while True: try: re_first_epoch.append(next(iterator)) except StopIteration: break assert re_first_epoch == all_rank_first_left_data[rank] # 查看第二轮的结果是否也是和第一次的第二轮完全一致； each_rank_iterable.set_epoch(1) iterator = iter(each_rank_iterable) re_second_epoch = [] for _ in range(forward_steps): re_second_epoch.append(next(iterator)) assert re_second_epoch == all_rank_second_epoch[rank] # todo 测试 ddp 时 world_size 改变的断点重训； def test_5(self): ...