chendl/compositional_test

1# coding=utf-82# Copyright 2018 the HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16import copy17import unittest18 19import numpy as np20 21from transformers.data.data_collator import default_data_collator22from transformers.testing_utils import require_accelerate, require_torch23from transformers.trainer_utils import RemoveColumnsCollator, find_executable_batch_size24from transformers.utils import is_torch_available25 26 27if is_torch_available():28    import torch29    from torch import nn30    from torch.utils.data import IterableDataset31 32    from transformers.modeling_outputs import SequenceClassifierOutput33    from transformers.tokenization_utils_base import BatchEncoding34    from transformers.trainer_pt_utils import (35        DistributedLengthGroupedSampler,36        DistributedSamplerWithLoop,37        DistributedTensorGatherer,38        IterableDatasetShard,39        LabelSmoother,40        LengthGroupedSampler,41        SequentialDistributedSampler,42        ShardSampler,43        get_parameter_names,44        numpy_pad_and_concatenate,45        torch_pad_and_concatenate,46    )47 48    class TstLayer(nn.Module):49        def __init__(self, hidden_size):50            super().__init__()51            self.linear1 = nn.Linear(hidden_size, hidden_size)52            self.ln1 = nn.LayerNorm(hidden_size)53            self.linear2 = nn.Linear(hidden_size, hidden_size)54            self.ln2 = nn.LayerNorm(hidden_size)55            self.bias = nn.Parameter(torch.zeros(hidden_size))56 57        def forward(self, x):58            h = self.ln1(nn.functional.relu(self.linear1(x)))59            h = nn.functional.relu(self.linear2(x))60            return self.ln2(x + h + self.bias)61 62    class RandomIterableDataset(IterableDataset):63        # For testing, an iterable dataset of random length64        def __init__(self, p_stop=0.01, max_length=1000):65            self.p_stop = p_stop66            self.max_length = max_length67            self.generator = torch.Generator()68 69        def __iter__(self):70            count = 071            stop = False72            while not stop and count < self.max_length:73                yield count74                count += 175                number = torch.rand(1, generator=self.generator).item()76                stop = number < self.p_stop77 78 79@require_torch80class TrainerUtilsTest(unittest.TestCase):81    def test_distributed_tensor_gatherer(self):82        # Simulate a result with a dataset of size 21, 4 processes and chunks of lengths 2, 3, 183        world_size = 484        num_samples = 2185        input_indices = [86            [0, 1, 6, 7, 12, 13, 18, 19],87            [2, 3, 4, 8, 9, 10, 14, 15, 16, 20, 0, 1],88            [5, 11, 17, 2],89        ]90 91        predictions = np.random.normal(size=(num_samples, 13))92        gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)93        for indices in input_indices:94            gatherer.add_arrays(predictions[indices])95        result = gatherer.finalize()96        self.assertTrue(np.array_equal(result, predictions))97 98        # With nested tensors99        gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)100        for indices in input_indices:101            gatherer.add_arrays([predictions[indices], [predictions[indices], predictions[indices]]])102        result = gatherer.finalize()103        self.assertTrue(isinstance(result, list))104        self.assertEqual(len(result), 2)105        self.assertTrue(isinstance(result[1], list))106        self.assertEqual(len(result[1]), 2)107        self.assertTrue(np.array_equal(result[0], predictions))108        self.assertTrue(np.array_equal(result[1][0], predictions))109        self.assertTrue(np.array_equal(result[1][1], predictions))110 111    def test_distributed_tensor_gatherer_different_shapes(self):112        # Simulate a result with a dataset of size 21, 4 processes and chunks of lengths 2, 3, 1113        world_size = 4114        num_samples = 21115        input_indices = [116            [0, 1, 6, 7, 12, 13, 18, 19],117            [2, 3, 4, 8, 9, 10, 14, 15, 16, 20, 0, 1],118            [5, 11, 17, 2],119        ]120        sequence_lengths = [8, 10, 13]121 122        predictions = np.random.normal(size=(num_samples, 13))123        gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)124        for indices, seq_length in zip(input_indices, sequence_lengths):125            gatherer.add_arrays(predictions[indices, :seq_length])126        result = gatherer.finalize()127 128        # Remove the extra samples added at the end for a round multiple of num processes.129        actual_indices = [input_indices[0], input_indices[1][:-2], input_indices[2][:-1]]130        for indices, seq_length in zip(actual_indices, sequence_lengths):131            self.assertTrue(np.array_equal(result[indices, :seq_length], predictions[indices, :seq_length]))132 133        # With nested tensors134        predictions = np.random.normal(size=(num_samples, 13))135        gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)136        for indices, seq_length in zip(input_indices, sequence_lengths):137            gatherer.add_arrays([predictions[indices, :seq_length], predictions[indices]])138        result = gatherer.finalize()139 140        for indices, seq_length in zip(actual_indices, sequence_lengths):141            self.assertTrue(np.array_equal(result[0][indices, :seq_length], predictions[indices, :seq_length]))142        self.assertTrue(np.array_equal(result[1], predictions))143 144        # Check if works if varying seq_length is second145        gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)146        for indices, seq_length in zip(input_indices, sequence_lengths):147            gatherer.add_arrays([predictions[indices], predictions[indices, :seq_length]])148        result = gatherer.finalize()149 150        self.assertTrue(np.array_equal(result[0], predictions))151        for indices, seq_length in zip(actual_indices, sequence_lengths):152            self.assertTrue(np.array_equal(result[1][indices, :seq_length], predictions[indices, :seq_length]))153 154    def test_label_smoothing(self):155        epsilon = 0.1156        num_labels = 12157        random_logits = torch.randn(4, 5, num_labels)158        random_labels = torch.randint(0, num_labels, (4, 5))159        loss = nn.functional.cross_entropy(random_logits.view(-1, num_labels), random_labels.view(-1))160        model_output = SequenceClassifierOutput(logits=random_logits)161        label_smoothed_loss = LabelSmoother(0.1)(model_output, random_labels)162        log_probs = -nn.functional.log_softmax(random_logits, dim=-1)163        expected_loss = (1 - epsilon) * loss + epsilon * log_probs.mean()164        self.assertTrue(torch.allclose(label_smoothed_loss, expected_loss))165 166        # With a few -100 labels167        random_labels[0, 1] = -100168        random_labels[2, 1] = -100169        random_labels[2, 3] = -100170 171        loss = nn.functional.cross_entropy(random_logits.view(-1, num_labels), random_labels.view(-1))172        model_output = SequenceClassifierOutput(logits=random_logits)173        label_smoothed_loss = LabelSmoother(0.1)(model_output, random_labels)174        log_probs = -nn.functional.log_softmax(random_logits, dim=-1)175        # Mask the log probs with the -100 labels176        log_probs[0, 1] = 0.0177        log_probs[2, 1] = 0.0178        log_probs[2, 3] = 0.0179        expected_loss = (1 - epsilon) * loss + epsilon * log_probs.sum() / (num_labels * 17)180        self.assertTrue(torch.allclose(label_smoothed_loss, expected_loss))181 182    def test_group_by_length(self):183        # Get some inputs of random lengths184        lengths = torch.randint(0, 25, (100,)).tolist()185        # Put one bigger than the others to check it ends up in first position186        lengths[32] = 50187 188        indices = list(LengthGroupedSampler(4, lengths=lengths))189        # The biggest element should be first190        self.assertEqual(lengths[indices[0]], 50)191        # The indices should be a permutation of range(100)192        self.assertEqual(sorted(indices), list(range(100)))193 194    def test_group_by_length_with_dict(self):195        # Get some inputs of random lengths196        data = []197        for _ in range(6):198            input_ids = torch.randint(0, 25, (100,)).tolist()199            data.append({"input_ids": input_ids})200        # Put one bigger than the others to check it ends up in first position201        data[3]["input_ids"] = torch.randint(0, 25, (105,)).tolist()202 203        indices = list(LengthGroupedSampler(4, dataset=data))204        # The biggest element should be first205        self.assertEqual(len(data[indices[0]]["input_ids"]), 105)206        # The indices should be a permutation of range(6)207        self.assertEqual(sorted(indices), list(range(6)))208 209    def test_group_by_length_with_batch_encoding(self):210        # Get some inputs of random lengths211        data = []212        for _ in range(6):213            input_ids = torch.randint(0, 25, (100,)).tolist()214            data.append(BatchEncoding({"input_ids": input_ids}))215        # Put one bigger than the others to check it ends up in first position216        data[3]["input_ids"] = torch.randint(0, 25, (105,)).tolist()217 218        indices = list(LengthGroupedSampler(4, dataset=data))219        # The biggest element should be first220        self.assertEqual(len(data[indices[0]]["input_ids"]), 105)221        # The indices should be a permutation of range(6)222        self.assertEqual(sorted(indices), list(range(6)))223 224    def test_distributed_length_grouped(self):225        # Get some inputs of random lengths226        lengths = torch.randint(0, 25, (100,)).tolist()227        # Put one bigger than the others to check it ends up in first position228        lengths[32] = 50229 230        indices_process_0 = list(DistributedLengthGroupedSampler(4, num_replicas=2, rank=0, lengths=lengths))231        indices_process_1 = list(DistributedLengthGroupedSampler(4, num_replicas=2, rank=1, lengths=lengths))232        # The biggest element should be first233        self.assertEqual(lengths[indices_process_0[0]], 50)234        # The indices should be a permutation of range(100)235        self.assertEqual(sorted(indices_process_0 + indices_process_1), list(range(100)))236 237    def test_get_parameter_names(self):238        model = nn.Sequential(TstLayer(128), nn.ModuleList([TstLayer(128), TstLayer(128)]))239        # fmt: off240        self.assertEqual(241            get_parameter_names(model, [nn.LayerNorm]),242            ['0.linear1.weight', '0.linear1.bias', '0.linear2.weight', '0.linear2.bias', '0.bias', '1.0.linear1.weight', '1.0.linear1.bias', '1.0.linear2.weight', '1.0.linear2.bias', '1.0.bias', '1.1.linear1.weight', '1.1.linear1.bias', '1.1.linear2.weight', '1.1.linear2.bias', '1.1.bias']243        )244        # fmt: on245 246    def test_distributed_sampler_with_loop(self):247        batch_size = 16248        for length in [23, 64, 123]:249            dataset = list(range(length))250            shard1 = DistributedSamplerWithLoop(dataset, batch_size, num_replicas=2, rank=0)251            shard2 = DistributedSamplerWithLoop(dataset, batch_size, num_replicas=2, rank=1)252 253            # Set seeds254            shard1.set_epoch(0)255            shard2.set_epoch(0)256 257            # Sample258            samples1 = list(shard1)259            samples2 = list(shard2)260 261            self.assertTrue(len(samples1) % batch_size == 0)262            self.assertTrue(len(samples2) % batch_size == 0)263 264            total = []265            for sample1, sample2 in zip(samples1, samples2):266                total += [sample1, sample2]267 268            self.assertEqual(set(total[:length]), set(dataset))269            self.assertEqual(set(total[length:]), set(total[: (len(total) - length)]))270 271    def test_sequential_distributed_sampler(self):272        batch_size = 16273        for length in [23, 64, 123]:274            dataset = list(range(length))275            shard1 = SequentialDistributedSampler(dataset, num_replicas=2, rank=0)276            shard2 = SequentialDistributedSampler(dataset, num_replicas=2, rank=1)277 278            # Sample279            samples1 = list(shard1)280            samples2 = list(shard2)281 282            total = samples1 + samples2283 284            self.assertListEqual(total[:length], dataset)285            self.assertListEqual(total[length:], dataset[: (len(total) - length)])286 287            # With a batch_size passed288            shard1 = SequentialDistributedSampler(dataset, num_replicas=2, rank=0, batch_size=batch_size)289            shard2 = SequentialDistributedSampler(dataset, num_replicas=2, rank=1, batch_size=batch_size)290 291            # Sample292            samples1 = list(shard1)293            samples2 = list(shard2)294 295            self.assertTrue(len(samples1) % batch_size == 0)296            self.assertTrue(len(samples2) % batch_size == 0)297 298            total = samples1 + samples2299 300            self.assertListEqual(total[:length], dataset)301            self.assertListEqual(total[length:], dataset[: (len(total) - length)])302 303    def check_iterable_dataset_shard(self, dataset, batch_size, drop_last, num_processes=2, epoch=0):304        # Set the seed for the base dataset to get the proper reference.305        dataset.generator.manual_seed(epoch)306        reference = list(dataset)307 308        shards = [309            IterableDatasetShard(310                dataset, batch_size=batch_size, drop_last=drop_last, num_processes=num_processes, process_index=i311            )312            for i in range(num_processes)313        ]314        for shard in shards:315            shard.set_epoch(epoch)316        shard_lists = [list(shard) for shard in shards]317 318        for shard in shard_lists:319            # All shards have a number of samples that is a round multiple of batch size320            self.assertTrue(len(shard) % batch_size == 0)321            # All shards have the same number of samples322            self.assertEqual(len(shard), len(shard_lists[0]))323 324        for shard in shards:325            # All shards know the total number of samples326            self.assertEqual(shard.num_examples, len(reference))327 328        observed = []329        for idx in range(0, len(shard_lists[0]), batch_size):330            for shard in shard_lists:331                observed += shard[idx : idx + batch_size]332 333        # If drop_last is False we loop through samples at the beginning to have a size that is a round multiple of334        # batch_size335        if not drop_last:336            while len(reference) < len(observed):337                reference += reference338        self.assertListEqual(observed, reference[: len(observed)])339 340        # Check equivalence between IterableDataset and ShardSampler341        dataset.generator.manual_seed(epoch)342        reference = list(dataset)343 344        sampler_shards = [345            ShardSampler(346                reference, batch_size=batch_size, drop_last=drop_last, num_processes=num_processes, process_index=i347            )348            for i in range(num_processes)349        ]350        for shard, sampler_shard in zip(shard_lists, sampler_shards):351            self.assertListEqual(shard, list(sampler_shard))352 353    def test_iterable_dataset_shard(self):354        dataset = RandomIterableDataset()355 356        self.check_iterable_dataset_shard(dataset, 4, drop_last=True, num_processes=2, epoch=0)357        self.check_iterable_dataset_shard(dataset, 4, drop_last=False, num_processes=2, epoch=0)358 359        self.check_iterable_dataset_shard(dataset, 4, drop_last=True, num_processes=3, epoch=42)360        self.check_iterable_dataset_shard(dataset, 4, drop_last=False, num_processes=3, epoch=42)361 362    def test_iterable_dataset_shard_with_length(self):363        sampler_shards = [364            IterableDatasetShard(list(range(100)), batch_size=4, drop_last=True, num_processes=2, process_index=i)365            for i in range(2)366        ]367 368        # Build expected shards: each process will have batches of size 4 until there is not enough elements to369        # form two full batches (so we stop at 96 = (100 // (4 * 2)) * 4)370        expected_shards = [[], []]371        current_shard = 0372        for i in range(0, 96, 4):373            expected_shards[current_shard].extend(list(range(i, i + 4)))374            current_shard = 1 - current_shard375 376        self.assertListEqual([list(shard) for shard in sampler_shards], expected_shards)377        self.assertListEqual([len(shard) for shard in sampler_shards], [len(shard) for shard in expected_shards])378 379        sampler_shards = [380            IterableDatasetShard(list(range(100)), batch_size=4, drop_last=False, num_processes=2, process_index=i)381            for i in range(2)382        ]383        # When drop_last=False, we get two last full batches by looping back to the beginning.384        expected_shards[0].extend(list(range(96, 100)))385        expected_shards[1].extend(list(range(0, 4)))386 387        self.assertListEqual([list(shard) for shard in sampler_shards], expected_shards)388        self.assertListEqual([len(shard) for shard in sampler_shards], [len(shard) for shard in expected_shards])389 390    def check_shard_sampler(self, dataset, batch_size, drop_last, num_processes=2):391        shards = [392            ShardSampler(393                dataset, batch_size=batch_size, drop_last=drop_last, num_processes=num_processes, process_index=i394            )395            for i in range(num_processes)396        ]397        shard_lists = [list(shard) for shard in shards]398 399        for shard in shard_lists:400            # All shards have a number of samples that is a round multiple of batch size401            self.assertTrue(len(shard) % batch_size == 0)402            # All shards have the same number of samples403            self.assertEqual(len(shard), len(shard_lists[0]))404 405        observed = []406        for idx in range(0, len(shard_lists[0]), batch_size):407            for shard in shard_lists:408                observed += shard[idx : idx + batch_size]409 410        # If drop_last is False we loop through samples at the beginning to have a size that is a round multiple of411        # batch_size412        reference = copy.copy(dataset)413        if not drop_last:414            while len(reference) < len(observed):415                reference += reference416        self.assertListEqual(observed, reference[: len(observed)])417 418    def test_shard_sampler(self):419        for n_elements in [64, 123]:420            dataset = list(range(n_elements))421 422            self.check_shard_sampler(dataset, 4, drop_last=True, num_processes=2)423            self.check_shard_sampler(dataset, 4, drop_last=False, num_processes=2)424 425            self.check_shard_sampler(dataset, 4, drop_last=True, num_processes=3)426            self.check_shard_sampler(dataset, 4, drop_last=False, num_processes=3)427 428    @require_accelerate429    def test_executable_batch_size(self):430        batch_sizes = []431 432        @find_executable_batch_size(starting_batch_size=64, auto_find_batch_size=True)433        def mock_training_loop_function(batch_size):434            nonlocal batch_sizes435            batch_sizes.append(batch_size)436            if batch_size > 16:437                raise RuntimeError("CUDA out of memory.")438 439        mock_training_loop_function()440        self.assertEqual(batch_sizes, [64, 32, 16])441 442    @require_accelerate443    def test_executable_batch_size_no_search(self):444        batch_sizes = []445 446        @find_executable_batch_size(starting_batch_size=64, auto_find_batch_size=False)447        def mock_training_loop_function(batch_size):448            nonlocal batch_sizes449            batch_sizes.append(batch_size)450 451        mock_training_loop_function()452        self.assertEqual(batch_sizes, [64])453 454    @require_accelerate455    def test_executable_batch_size_with_error(self):456        @find_executable_batch_size(starting_batch_size=64, auto_find_batch_size=False)457        def mock_training_loop_function(batch_size):458            raise RuntimeError("CUDA out of memory.")459 460        with self.assertRaises(RuntimeError) as cm:461            mock_training_loop_function()462            self.assertEqual("CUDA out of memory", cm.args[0])463 464    def test_pad_and_concatenate_with_1d(self):465        """Tests whether pad_and_concatenate works with scalars."""466        array1 = 1.0467        array2 = 2.0468        result = numpy_pad_and_concatenate(array1, array2)469        self.assertTrue(np.array_equal(np.array([1.0, 2.0]), result))470 471        tensor1 = torch.tensor(1.0)472        tensor2 = torch.tensor(2.0)473        result = torch_pad_and_concatenate(tensor1, tensor2)474        self.assertTrue(torch.equal(result, torch.Tensor([1.0, 2.0])))475 476    def test_remove_columns_collator(self):477        class MockLogger:478            def __init__(self) -> None:479                self.called = 0480 481            def info(self, msg):482                self.called += 1483                self.last_msg = msg484 485        data_batch = [486            {"col1": 1, "col2": 2, "col3": 3},487            {"col1": 1, "col2": 2, "col3": 3},488        ]489        logger = MockLogger()490        remove_columns_collator = RemoveColumnsCollator(491            default_data_collator, ["col1", "col2"], logger, "model", "training"492        )493 494        self.assertNotIn("col3", remove_columns_collator(data_batch))495        # check that the logging message is printed out only once496        remove_columns_collator(data_batch)497        remove_columns_collator(data_batch)498        self.assertEqual(logger.called, 1)499        self.assertIn("col3", logger.last_msg)500